diff --git a/agent.py b/agent.py
index 43c998d..7554f64 100644
--- a/agent.py
+++ b/agent.py
@@ -120,6 +120,7 @@ class QuoteJudgeStudioV1(A2AAgent[QuoteJudgeStudioV1Config, PlatformUserAuth]):
     config_model = QuoteJudgeStudioV1Config
     auth_model = PlatformUserAuth
     state = State.DURABLE
+    state_model = QuoteJudgeState
     resources = Resources(cpu="500m", memory="512Mi", max_runtime_seconds=300)
     platform_resources = AgentPlatformResources(
         databases=(