DEV Community

Yash Lohade
Yash Lohade

Posted on

How I Debugged My AI Trading Agent Token Overflow, Memory Leaks & Silent NaN Predictions

Summer Bug Smash: Clear the Lineup 🐛🛹

Project Overview

Algomaya is an AI-powered algorithmic trading education platform I'm building solo. It lets users learn algo trading, build strategies with a no-code builder, backtest against historical data, and practice with paper trading — all powered by an AI agent called Maya that autonomously chains together tools (backtests, signals, quotes, news) to answer complex trading questions.

Stack: React Native (Expo SDK 54) + FastAPI + PostgreSQL + Groq LLM (Llama 3.3 70B) + LSTM neural network for price prediction

The app has 89+ screens, 335+ commits, and a production backend on Google Cloud Run. It's a real product heading toward Play Store launch — which is exactly why these bugs mattered.


Bug Fix or Performance Improvement

I fixed 6 critical bugs across the AI agent pipeline, LSTM prediction model, and frontend — the kind of bugs that silently corrupt data, leak memory, and crash production.

Bug 1: AI Agent Token Overflow — Context Window Bomb

File: backend/app/services/maya_agent.py

Maya's ReAct loop accumulates messages with every LLM call and tool result. Each tool response is truncated to 2000 chars, but with 15 max steps, the message array can balloon past Groq's context window limit. The result: 400 Bad Request errors that kill the agent mid-analysis, or worse — silently truncated context that makes the LLM hallucinate.

The fix — sliding window token management:

# BEFORE: Messages grew unbounded
messages: List[dict] = [
    {"role": "system", "content": MAYA_SYSTEM_PROMPT},
    {"role": "user", "content": goal},
]
# ... messages.append() on every iteration, no limit check

# AFTER: Token-aware trimming keeps context within budget
@staticmethod
def _estimate_token_count(messages: List[dict]) -> int:
    """Rough token estimate (~1.3 tokens per word)."""
    total = 0
    for m in messages:
        content = m.get("content", "")
        if isinstance(content, str):
            total += len(content.split())
    return int(total * 1.3)

@staticmethod
def _trim_messages(messages: List[dict], max_tokens: int = 6000) -> List[dict]:
    """Keep system prompt + user goal + most recent messages within budget."""
    if MayaAgent._estimate_token_count(messages) <= max_tokens:
        return messages
    # Always preserve system prompt (index 0) and user goal (index 1)
    trimmed = [messages[0], messages[1]]
    remaining = messages[2:]
    tail: List[dict] = []
    budget = max_tokens - MayaAgent._estimate_token_count(trimmed)
    for msg in reversed(remaining):
        msg_tokens = int(len(str(msg.get("content", "")).split()) * 1.3)
        if budget - msg_tokens < 0:
            break
        tail.insert(0, msg)
        budget -= msg_tokens
    return trimmed + tail
Enter fullscreen mode Exit fullscreen mode

Impact: Prevents context window overflow on complex multi-tool agent runs. The agent now gracefully drops older tool results while always preserving the system prompt and user's original goal.


Bug 2: Agent Run Hangs Indefinitely — No Outer Timeout

File: backend/app/services/maya_agent.py

Each Groq API call has a 30s timeout, but the entire agent loop had no outer timeout. If the LLM kept requesting tools in a loop (tool → LLM → tool → LLM...), the backend task could run for minutes or hours. The frontend times out at 30s, but the backend keeps burning CPU and API credits.

# BEFORE: No timeout guard on the loop
while step_number < max_steps:
    response = await _call_groq(messages, tools=MAYA_TOOLS)
    # ... could run forever if LLM keeps calling tools

# AFTER: Hard 5-minute timeout
MAX_RUN_TIMEOUT_SECONDS = 300

while step_number < max_steps:
    elapsed = (datetime.now(timezone.utc) - start_time).total_seconds()
    if elapsed > MAX_RUN_TIMEOUT_SECONDS:
        logger.warning(f"Maya run {run_id} exceeded {MAX_RUN_TIMEOUT_SECONDS}s timeout")
        run.summary = "Analysis timed out. Please try a simpler query."
        run.status = "completed"
        db.commit()
        return
    # ... rest of loop
Enter fullscreen mode Exit fullscreen mode

Impact: Prevents runaway agent executions from burning Groq API credits and Cloud Run CPU. Users get a clean timeout message instead of an infinite spinner.


Bug 3: Tool Result Truncation Destroys LLM Context

File: backend/app/services/maya_agent.py

When a backtest returns 500 trades, the old _truncate_result() replaced the entire array with the string "[500 items omitted]". The LLM then received this useless string instead of actual data, making it impossible to analyze results properly.

# BEFORE: Replace arrays with useless strings
if k in ("chart_data", "trades", "equity_curve"):
    truncated[k] = f"[{len(v)} items omitted]"  # LLM can't parse this

# AFTER: Keep first + last N items so LLM has real data to reason over
if k == "trades":
    if isinstance(v, list) and len(v) > max_items:
        truncated[k] = v[:5] + v[-5:]  # First 5 + last 5 trades
        truncated["_trades_total"] = len(v)  # Total count for context
Enter fullscreen mode Exit fullscreen mode

Impact: The LLM can now actually analyze backtest results — seeing both early and recent trades, real equity curve data points, and accurate totals. Before this fix, Maya would often say "I ran the backtest but can't see the results" because the data was replaced with placeholder strings.


Bug 4: LSTM Model Silently Predicts NaN — No Scaler Guards

File: backend/app/services/lstm_model.py

The LSTM price prediction model's _scale_sequences() and predict() methods had zero protection against NaN/Inf values. When market data had gaps (common with Indian stocks on holidays), the scaler would produce NaN, the model would predict NaN, and the frontend would silently show "NEUTRAL" for every stock — with no error logged anywhere.

# BEFORE: No NaN protection — silent data corruption
def _scale_sequences(self, X_seq):
    flat = X_seq.reshape(-1, n_features)
    scaled = self.scaler.transform(flat)  # NaN in → NaN out, silently
    return scaled.reshape(n_samples, seq_len, n_features)

# AFTER: Detect, log, and handle NaN at every stage
def _scale_sequences(self, X_seq):
    if self.scaler is None:
        raise ValueError("Scaler not fitted. Call train() or load() first.")

    flat = X_seq.reshape(-1, n_features)

    if np.isnan(flat).any():
        logger.warning("Input contains NaN values before scaling — forward-filling")
        flat = pd.DataFrame(flat).ffill().fillna(0).values

    scaled = self.scaler.transform(flat)

    if np.isnan(scaled).any() or np.isinf(scaled).any():
        logger.warning("Scaling produced NaN/Inf — clamping to safe range")
        scaled = np.nan_to_num(scaled, nan=0.0, posinf=3.0, neginf=-3.0)

    return scaled.reshape(n_samples, seq_len, n_features)
Enter fullscreen mode Exit fullscreen mode

Impact: Stock predictions no longer silently fail. NaN values from missing market data are forward-filled (a standard financial data technique), and extreme values that break the scaler are clamped to a safe range. Users now see actual predictions instead of blanket "NEUTRAL" signals.


Bug 5: TensorFlow Memory Leak in Prediction Loop

File: backend/app/services/lstm_model.py

Every call to predict() allocated X_scaled and X_seq numpy arrays but never cleaned them up. On Cloud Run with limited memory, repeated predictions (e.g., scanning 50 stocks for signals) would accumulate hundreds of MB of orphaned arrays until the container OOM-killed.

# BEFORE: Arrays allocated, never freed
X_scaled = self.scaler.transform(X)
X_seq = self.create_sequences_predict(X_scaled)
probabilities = self.model.predict(X_seq, verbose=0).flatten()
# X_scaled and X_seq still in memory

# AFTER: Explicit cleanup in finally block
try:
    probabilities = self.model.predict(X_seq, verbose=0).flatten()
    predictions = (probabilities >= 0.5).astype(int)
finally:
    del X_scaled, X_seq
    import gc
    gc.collect()
Enter fullscreen mode Exit fullscreen mode

Impact: Prevents OOM crashes on Cloud Run during batch prediction operations. Memory is explicitly freed after each prediction call, keeping the container's footprint stable even when scanning dozens of stocks.


Bug 6: Sentry AI Agent Monitoring Integration

File: backend/app/services/maya_agent.py, backend/app/main.py, backend/app/core/config.py

Beyond fixing bugs, I integrated Sentry for full AI agent observability. Each Maya agent run now creates a Sentry transaction with spans for every LLM call and tool execution, complete with token usage tracking and error capture.

# Sentry transaction wraps the entire agent run
with sentry_sdk.start_transaction(
    op="ai.agent", name="maya_agent_run", description=goal[:100]
) as txn:
    txn.set_tag("tier", tier)
    txn.set_tag("run_id", str(run_id))

    # Each LLM call gets its own span with token tracking
    with sentry_sdk.start_span(
        op="ai.chat_completions.create", description=f"groq/{GROQ_MODEL}"
    ) as llm_span:
        llm_span.set_data("ai.model_id", GROQ_MODEL)
        response = await _call_groq(messages, tools=MAYA_TOOLS)

        usage = response.get("usage", {})
        llm_span.set_data("ai.prompt_tokens_used", usage.get("prompt_tokens", 0))
        llm_span.set_data("ai.completion_tokens_used", usage.get("completion_tokens", 0))

    # Each tool execution gets a span
    with sentry_sdk.start_span(
        op="ai.tool", description=f"tool:{tool_name}"
    ) as tool_span:
        tool_span.set_data("ai.tool.name", tool_name)
        result = await _execute_tool(tool_name, tool_args, tier=tier)
Enter fullscreen mode Exit fullscreen mode

What this enables:

  • Conversation traces — See the full agent reasoning chain (LLM call → tool → LLM → tool → summary) as a waterfall in Sentry
  • Token usage monitoring — Track prompt/completion tokens per run to detect context overflow before it happens
  • Tool failure tracking — Pinpoint which tools fail most often and why
  • Performance profiling — Identify slow tool executions (backtests taking 10s+) vs fast ones (quotes <1s)

Code

All code changes are in the diffs above. Summary of files modified:

File Change Impact
backend/app/services/maya_agent.py Token management, timeout guard, truncation fix, Sentry AI tracing Prevents crashes, hangs, and data loss in AI agent
backend/app/services/lstm_model.py NaN guards, memory cleanup Prevents silent prediction failures and OOM crashes
backend/app/main.py Sentry SDK initialization Error tracking and performance monitoring
backend/app/core/config.py Sentry config variables Environment-specific monitoring config
backend/requirements.txt Addedsentry-sdk[fastapi] Sentry dependency

My Improvements

These fixes address the hardest class of bugs in AI applications — the ones that don't crash loudly but silently corrupt your data pipeline:

  1. Token overflow was making Maya hallucinate on complex queries because the LLM lost context. The sliding window approach keeps the most recent tool results while always preserving the original goal.
  2. The missing timeout was burning real money — each Groq API call costs tokens, and a runaway agent could chain 15 calls over several minutes with no kill switch.
  3. NaN propagation is the classic ML pipeline bug. One missing data point in market data → NaN features → NaN scaled values → NaN predictions → "NEUTRAL" everywhere. The fix adds guards at every stage of the pipeline.
  4. The Sentry integration isn't just monitoring — it's AI agent observability. Each agent run becomes a trace with spans for LLM calls and tool executions, token usage metrics, and error capture. This is exactly what you need to debug why an AI agent gave a bad answer.

How I found them: I conducted a systematic audit documenting 71 bugs across 4 severity levels. The AI pipeline bugs were the most dangerous because they failed silently — no crash, no error log, just wrong results served to users.


Best Use of Sentry

I integrated Sentry (sentry-sdk[fastapi]>=2.10.0) into my FastAPI backend specifically for AI agent monitoring:

Setup:

  • Added SENTRY_DSN, SENTRY_TRACES_SAMPLE_RATE, and SENTRY_PROFILES_SAMPLE_RATE to my Pydantic settings config
  • Initialized Sentry in main.py before middleware with enable_tracing=True
  • Used Sentry's AI-specific span operations (ai.agent, ai.chat_completions.create, ai.tool) for proper categorization

How Sentry helps debug my AI agent:

  1. Conversation Traces — Each Maya agent run creates a Sentry transaction. I can see the full reasoning chain as a waterfall: which tools were called, how long each LLM call took, and where failures occurred.
  2. Token Usage Tracking — Every LLM span records ai.prompt_tokens_used and ai.completion_tokens_used. I can now detect token overflow before it hits the context limit by monitoring token growth across runs.
  3. Tool Failure Monitoring — When a tool fails (e.g., Yahoo Finance API timeout), Sentry captures the exception with full context: which stock, which strategy, what the LLM was trying to do. This lets me prioritize which tool integrations need hardening.
  4. Performance Profiling — Sentry's continuous profiling shows that backtests are the bottleneck (5-15s per run), while stock quotes are fast (<500ms). This informed my decision to add the timeout guard — a 15-step run with 10 backtests could easily exceed 2 minutes.
  5. Error Alerting — Sentry alerts on new error types. The first time I deployed, it immediately caught a json.JSONDecodeError from malformed LLM tool arguments that I'd been silently swallowing as empty {} — a bug I didn't even know existed.

Top comments (0)