diff --git a/agent.py b/agent.py
index 07c5ff5..1f495b2 100644
--- a/agent.py
+++ b/agent.py
@@ -126,6 +126,7 @@ class ContractClockStudioV1(A2AAgent[ContractClockStudioV1Config, PlatformUserAu
         max_total_size_bytes=MAX_BROWSER_UPLOAD_BYTES,
     )
     state = State.DURABLE
+    state_model = ContractClockState
     tools_used = ("managed-postgres", "mcp", "packed-frontend")
     platform_resources = AgentPlatformResources(
         databases=(