From 577d214f45c9895bf58da1d9efe9430e335b809e Mon Sep 17 00:00:00 2001
From: Nicholas Eckardt <neckardt@amazon.com>
Date: Wed, 25 Feb 2026 21:23:05 -0800
Subject: [PATCH] feat(agent): add generic type parameter to
 AgentState.getAll()

The get, set, and delete methods all accept a <TState> generic for
type-safe state access, but getAll() forced callers to use `as` casts.
This adds a generic overload so callers can opt into a narrower return
type (e.g. state.getAll<AppState>()) without a cast, while the
unparameterized overload continues to return Record<string, JSONValue>.

No runtime behavior changes.
---
 src/agent/__tests__/state.test.ts | 13 +++++++++++++
 src/agent/state.ts                | 12 ++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/src/agent/__tests__/state.test.ts b/src/agent/__tests__/state.test.ts
index 73ffccab..5df356e1 100644
--- a/src/agent/__tests__/state.test.ts
+++ b/src/agent/__tests__/state.test.ts
@@ -302,6 +302,19 @@ describe('AgentState', () => {
       const state = new AgentState()
       expect(state.getAll()).toEqual({})
     })
+
+    it('supports typed usage with generic state interface', () => {
+      interface MyState extends Record<string, number | string> {
+        count: number
+        name: string
+      }
+
+      const state = new AgentState({ count: 10, name: 'test' })
+      const result = state.getAll<MyState>()
+
+      expect(result.count).toBe(10)
+      expect(result.name).toBe('test')
+    })
   })
 
   describe('keys', () => {
diff --git a/src/agent/state.ts b/src/agent/state.ts
index ff5b313a..26fb5f06 100644
--- a/src/agent/state.ts
+++ b/src/agent/state.ts
@@ -124,8 +124,20 @@ export class AgentState implements StateSerializable {
   /**
    * Get a copy of all state as an object.
    *
+   * @typeParam TState - Optional state interface type for typed returns
    * @returns Deep copy of all state
+   *
+   * @example
+   * ```typescript
+   * // Typed usage
+   * const all = state.getAll<AppState>()   // AppState
+   *
+   * // Untyped usage
+   * const raw = state.getAll()             // Record<string, JSONValue>
+   * ```
    */
+  getAll<TState extends Record<string, JSONValue>>(): TState
+  getAll(): Record<string, JSONValue>
   getAll(): Record<string, JSONValue> {
     return deepCopy(this._state) as Record<string, JSONValue>
   }