1
// src/state/dynamo-state.ts — Agent state store backed by DynamoDB
2

3
import { DynamoDBClient } from "@aws-sdk/client-dynamodb";
4
import {
5
  DynamoDBDocumentClient,
6
  GetCommand,
7
  PutCommand,
8
  QueryCommand,
9
  DeleteCommand,
10
  BatchWriteCommand,
11
} from "@aws-sdk/lib-dynamodb";
12

13
export interface AgentSession {
14
  sessionId: string;
15
  conversation: Array<{
16
    role: "user" | "assistant" | "tool";
17
    content: string;
18
    timestamp: number;
19
    metadata?: Record<string, unknown>;
20
  }>;
21
  state: Record<string, unknown>;
22
  version: number;
23
  lastToolCall: string | null;
24
  lastActiveRegion: string;
25
  lastActiveAt: number;
26
}
27

28
export interface ToolCacheEntry {
29
  toolName: string;
30
  paramsHash: string;
31
  result: unknown;
32
  cachedAt: number;
33
  expiresAt: number;
34
  hitCount: number;
35
}
36

37
export class DynamoAgentState {
38
  private docClient: DynamoDBDocumentClient;
39
  private sessionTable: string;
40
  private cacheTable: string;
41

42
  constructor(options: {
43
    region: string;
44
    sessionTable?: string;
45
    cacheTable?: string;
46
  }) {
47
    const client = new DynamoDBClient({ region: options.region });
48
    this.docClient = DynamoDBDocumentClient.from(client);
49
    this.sessionTable = options.sessionTable || "agent-sessions";
50
    this.cacheTable = options.cacheTable || "agent-tool-cache";
51
  }
52

53
  // ──── Session Operations ────
54

55
  /**
56
   * Save or update a session checkpoint.
57
   * Uses conditional write for version-wins conflict resolution.
58
   */
59
  async saveSession(session: AgentSession): Promise<void> {
60
    const ttl = Math.floor(Date.now() / 1000) + 3600; // 1 hour TTL
61

62
    const item = {
63
      pk: `session#${session.sessionId}`,
64
      sk: "meta",
65
      data: {
66
        conversation: session.conversation.slice(-50), // Keep last 50 messages
67
        state: session.state,
68
      },
69
      version: session.version,
70
      lastToolCall: session.lastToolCall,
71
      lastActiveRegion: session.lastActiveRegion,
72
      lastActiveAt: session.lastActiveAt,
73
      ttl,
74
    };
75

76
    await this.docClient.send(new PutCommand({
77
      TableName: this.sessionTable,
78
      Item: item,
79
      // Only write if version is higher (prevent stale writes)
80
      ConditionExpression: "attribute_not_exists(version) OR version < :newVersion",
81
      ExpressionAttributeValues: {
82
        ":newVersion": session.version,
83
      },
84
    }));
85
  }
86

87
  /**
88
   * Load a session from DynamoDB.
89
   */
90
  async loadSession(sessionId: string): Promise<AgentSession | null> {
91
    const result = await this.docClient.send(new GetCommand({
92
      TableName: this.sessionTable,
93
      Key: {
94
        pk: `session#${sessionId}`,
95
        sk: "meta",
96
      },
97
      ConsistentRead: true, // Strong consistency for session reads
98
    }));
99

100
    if (!result.Item) return null;
101

102
    return this.hydrateSession(sessionId, result.Item);
103
  }
104

105
  /**
106
   * Append a single message to a session.
107
   * More efficient than rewriting the entire conversation.
108
   */
109
  async appendMessage(
110
    sessionId: string,
111
    message: AgentSession["conversation"][0]
112
  ): Promise<void> {
113
    const msgId = `msg#${String(message.timestamp).padStart(14, "0")}`;
114
    const ttl = Math.floor(Date.now() / 1000) + 3600;
115

116
    await this.docClient.send(new PutCommand({
117
      TableName: this.sessionTable,
118
      Item: {
119
        pk: `session#${sessionId}`,
120
        sk: msgId,
121
        role: message.role,
122
        content: message.content,
123
        timestamp: message.timestamp,
124
        metadata: message.metadata,
125
        ttl,
126
      },
127
    }));
128
  }
129

130
  /**
131
   * Load full conversation (meta + all messages).
132
   */
133
  async loadConversation(sessionId: string): Promise<AgentSession["conversation"]> {
134
    const result = await this.docClient.send(new QueryCommand({
135
      TableName: this.sessionTable,
136
      KeyConditionExpression: "pk = :pk",
137
      ExpressionAttributeValues: {
138
        ":pk": `session#${sessionId}`,
139
      },
140
      ConsistentRead: true,
141
    }));
142

143
    if (!result.Items || result.Items.length === 0) return [];
144

145
    // Sort messages by timestamp and extract conversation
146
    const messages = result.Items
147
      .filter(item => item.sk !== "meta")
148
      .sort((a, b) => (a.timestamp || 0) - (b.timestamp || 0))
149
      .map(item => ({
150
        role: item.role as "user" | "assistant" | "tool",
151
        content: item.content as string,
152
        timestamp: item.timestamp as number,
153
        metadata: item.metadata as Record<string, unknown> | undefined,
154
      }));
155

156
    return messages;
157
  }
158

159
  /**
160
   * Delete a session (cleanup on session end).
161
   */
162
  async deleteSession(sessionId: string): Promise<void> {
163
    // Get all items for this session
164
    const items = await this.docClient.send(new QueryCommand({
165
      TableName: this.sessionTable,
166
      KeyConditionExpression: "pk = :pk",
167
      ExpressionAttributeValues: { ":pk": `session#${sessionId}` },
168
    }));
169

170
    if (!items.Items || items.Items.length === 0) return;
171

172
    // Batch delete all items
173
    const deleteRequests = items.Items.map(item => ({
174
      DeleteRequest: {
175
        Key: { pk: item.pk, sk: item.sk },
176
      },
177
    }));
178

179
    // Batch in chunks of 25 (DynamoDB limit)
180
    for (let i = 0; i < deleteRequests.length; i += 25) {
181
      await this.docClient.send(new BatchWriteCommand({
182
        RequestItems: {
183
          [this.sessionTable]: deleteRequests.slice(i, i + 25),
184
        },
185
      }));
186
    }
187
  }
188

189
  // ──── Tool Cache Operations ────
190

191
  /**
192
   * Store a tool result in the cache table.
193
   */
194
  async cacheToolResult(
195
    toolName: string,
196
    params: Record<string, unknown>,
197
    result: unknown,
198
    ttlMs: number
199
  ): Promise<void> {
200
    const paramsHash = this.hashParams(params);
201
    const now = Date.now();
202
    const ttl = Math.floor((now + ttlMs) / 1000);
203

204
    await this.docClient.send(new PutCommand({
205
      TableName: this.cacheTable,
206
      Item: {
207
        pk: `tool#${toolName}#${paramsHash}`,
208
        result: typeof result === "string" ? result : JSON.stringify(result),
209
        cachedAt: now,
210
        expiresAt: now + ttlMs,
211
        hitCount: 0,
212
        ttl,
213
      },
214
    }));
215
  }
216

217
  /**
218
   * Get a cached tool result.
219
   */
220
  async getCachedToolResult<T>(
221
    toolName: string,
222
    params: Record<string, unknown>
223
  ): Promise<{ result: T; hitCount: number } | null> {
224
    const paramsHash = this.hashParams(params);
225

226
    const result = await this.docClient.send(new GetCommand({
227
      TableName: this.cacheTable,
228
      Key: { pk: `tool#${toolName}#${paramsHash}` },
229
    }));
230

231
    if (!result.Item) return null;
232

233
    // Update hit count (fire and forget)
234
    this.incrementCacheHit(toolName, paramsHash).catch(() => {});
235

236
    return {
237
      result: typeof result.Item.result === "string"
238
        ? JSON.parse(result.Item.result)
239
        : result.Item.result as T,
240
      hitCount: (result.Item.hitCount || 0) + 1,
241
    };
242
  }
243

244
  /**
245
   * Invalidate cache for a tool pattern.
246
   */
247
  async invalidateCache(toolName: string, paramsHash?: string): Promise<void> {
248
    if (paramsHash) {
249
      // Delete specific entry
250
      await this.docClient.send(new DeleteCommand({
251
        TableName: this.cacheTable,
252
        Key: { pk: `tool#${toolName}#${paramsHash}` },
253
      }));
254
    } else {
255
      // Query all entries for this tool and delete
256
      const items = await this.docClient.send(new QueryCommand({
257
        TableName: this.cacheTable,
258
        KeyConditionExpression: "begins_with(pk, :prefix)",
259
        ExpressionAttributeValues: { ":prefix": `tool#${toolName}#` },
260
      }));
261

262
      if (items.Items && items.Items.length > 0) {
263
        const deletes = items.Items.map(item => ({
264
          DeleteRequest: { Key: { pk: item.pk } },
265
        }));
266
        for (let i = 0; i < deletes.length; i += 25) {
267
          await this.docClient.send(new BatchWriteCommand({
268
            RequestItems: { [this.cacheTable]: deletes.slice(i, i + 25) },
269
          }));
270
        }
271
      }
272
    }
273
  }
274

275
  // ──── Admin / Stats ────
276

277
  /**
278
   * Get table statistics.
279
   */
280
  async getStats(): Promise<{
281
    activeSessions: number;
282
    cacheEntries: number;
283
  }> {
284
    const [sessionInfo, cacheInfo] = await Promise.all([
285
      this.docClient.send(new GetCommand({
286
        TableName: this.sessionTable,
287
        Key: { pk: "stats", sk: "meta" },
288
      })),
289
      this.docClient.send(new GetCommand({
290
        TableName: this.cacheTable,
291
        Key: { pk: "stats", sk: "meta" },
292
      })),
293
    ]);
294

295
    return {
296
      activeSessions: (sessionInfo.Item?.count as number) || 0,
297
      cacheEntries: (cacheInfo.Item?.count as number) || 0,
298
    };
299
  }
300

301
  // ──── Private ────
302

303
  private hydrateSession(sessionId: string, item: Record<string, any>): AgentSession {
304
    const data = item.data || {};
305
    return {
306
      sessionId,
307
      conversation: data.conversation || [],
308
      state: data.state || {},
309
      version: item.version || 1,
310
      lastToolCall: item.lastToolCall || null,
311
      lastActiveRegion: item.lastActiveRegion || "unknown",
312
      lastActiveAt: item.lastActiveAt || Date.now(),
313
    };
314
  }
315

316
  private hashParams(params: Record<string, unknown>): string {
317
    const crypto = require("crypto");
318
    return crypto
319
      .createHash("sha256")
320
      .update(JSON.stringify(params))
321
      .digest("hex")
322
      .slice(0, 16);
323
  }
324

325
  private async incrementCacheHit(
326
    toolName: string,
327
    paramsHash: string
328
  ): Promise<void> {
329
    await this.docClient.send(new PutCommand({
330
      TableName: this.cacheTable,
331
      Item: {
332
        pk: `tool#${toolName}#${paramsHash}`,
333
        hitCount: 1,
334
      },
335
      // Atomic increment using UpdateItem would be better
336
      // This is a simplified example
337
    }));
338
  }
339
}

1
import { DynamoDBDocumentClient } from "@aws-sdk/lib-dynamodb";
2
import { DynamoDBClient } from "@aws-sdk/client-dynamodb";
3

4
// Standard DynamoDB (for writes and cache misses)
5
const standardClient = DynamoDBDocumentClient.from(
6
  new DynamoDBClient({ region: "us-east-1" })
7
);
8

9
// With DAX — install: npm install amazon-dax-client
10
// import AmazonDaxClient from "amazon-dax-client";
11
// const daxClient = DynamoDBDocumentClient.from(
12
//   new AmazonDaxClient({ endpoints: ["agent-state-dax.xxxxx.clusters.dax.amazonaws.com"] })
13
// );
14

15
export class ReadOptimizedState extends DynamoAgentState {
16
  /**
17
   * Load session from DAX, fall back to DynamoDB.
18
   */
19
  async loadSessionFast(sessionId: string): Promise<AgentSession | null> {
20
    try {
21
      // Try DAX first
22
      // const result = await daxClient.send(...);
23
      // return result;
24
      return await this.loadSession(sessionId);
25
    } catch {
26
      return await this.loadSession(sessionId);
27
    }
28
  }
29
}