Race condition in InvocationReplayState.incrementAgentReplayIndex during parallel agent execution

[Harishvishwa01]

---

Description

InvocationReplayState maintains replay indices per agent using a HashMap:

private final Map<String, Integer> agentReplayIndices = new HashMap<>();

The class documentation mentions:

"Per-invocation replay state to isolate concurrent runs"
"Per-agent replay indices for parallel execution"

However, HashMap is not thread-safe, and the method incrementAgentReplayIndex performs a non-atomic read-modify-write sequence:

public void incrementAgentReplayIndex(String agentName) {
  int currentIndex = getAgentReplayIndex(agentName);
  setAgentReplayIndex(agentName, currentIndex + 1);
}

If multiple threads update the same agent index concurrently, the following race condition may occur:

1. Thread A reads index 1
2. Thread B reads index 1
3. Thread A writes 2
4. Thread B writes 2

The expected value should be 3, but the result becomes 2.

This can cause incorrect replay ordering or skipped increments during parallel agent execution.

---

Suggested Improvement

Use a thread-safe structure such as ConcurrentHashMap and perform atomic updates using merge().

Example improvement:

import java.util.concurrent.ConcurrentHashMap;

private final Map<String, Integer> agentReplayIndices = new ConcurrentHashMap<>();

public void incrementAgentReplayIndex(String agentName) {
  agentReplayIndices.merge(agentName, 1, Integer::sum);
}

This ensures the increment operation is atomic and prevents lost updates when multiple agents update the same index concurrently.

---

Additional Context

Since the class explicitly supports parallel execution, ensuring thread-safe replay index updates will make the replay mechanism deterministic and more reliable under concurrent agent runs.

[hemasekhar-p]

Hi @Harishvishwa01, thank you for reporting this issue. I have tried reproducing it using your inputs, but unfortunately, I was not able to replicate the behavior. To help us investigate and reproduce the problem, could you please share a minimal, reproducible example of your setup?

Specifically, it would be helpful if you could provide:

How your agents are executed?
How the ReplayPlugin is being triggered in your project?
Any configuration that enables replay functionality?

[Harishvishwa01]

Hi @hemasekhar-p , thanks for checking.

I didn’t run a full ADK project setup. I raised this issue based on reviewing the implementation of InvocationReplayState and ReplayPlugin, where I noticed a potential race condition pattern in the replay index update logic.

From the code comments mentioning parallel execution, I assumed agents may run concurrently and callbacks like beforeModelCallback or beforeToolCallback could potentially execute in parallel during an invocation.

Replay mode appears to be enabled when _adk_replay_config is present in the session state with:

• dir
• user_message_index

One clarification that would help: Is InvocationReplayState guaranteed to be accessed from a single thread per invocation?

• If yes → the current implementation should be safe.
• If not → an atomic increment (e.g., ConcurrentHashMap.merge) could help avoid potential race conditions.