diff --git a/core/src/main/java/com/google/adk/sessions/Session.java b/core/src/main/java/com/google/adk/sessions/Session.java
index f8376589a..94504fd96 100644
--- a/core/src/main/java/com/google/adk/sessions/Session.java
+++ b/core/src/main/java/com/google/adk/sessions/Session.java
@@ -27,8 +27,8 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ConcurrentMap;
 
 /** A {@link Session} object that encapsulates the {@link State} and {@link Event}s of a session. */
 @JsonDeserialize(builder = Session.Builder.class)
@@ -101,7 +101,7 @@ public Builder state(State state) {
 
     @CanIgnoreReturnValue
     @JsonProperty("state")
-    public Builder state(ConcurrentMap<String, Object> state) {
+    public Builder state(Map<String, Object> state) {
       this.state = new State(state);
       return this;
     }
@@ -162,7 +162,7 @@ public String id() {
   }
 
   @JsonProperty("state")
-  public ConcurrentMap<String, Object> state() {
+  public Map<String, Object> state() {
     return state;
   }