@aws/agentcore 0.21.1 → 0.23.0

package/dist/assets/typescript/http/vercelai/base/main.ts

@@ -1,22 +1,62 @@
 import { BedrockAgentCoreApp } from 'bedrock-agentcore/runtime';
-import { streamText } from 'ai';
+import { streamText, type ModelMessage } from 'ai';
 import { loadModel } from './model/load.js';
 
 const SYSTEM_PROMPT = `You are a helpful assistant.`;
 
+const HISTORY_LIMIT = 128;
+
+// Keeps one message history per sessionId so each session remembers its own
+// turns (best-effort; resets on cold start). A Map preserves insertion order,
+// so it doubles as an LRU bounded to 128 sessions — a local dev process serving
+// many sessions cannot leak history between them or grow without bound. On
+// AgentCore Runtime each microVM serves a single session, so this holds one
+// entry. For durable history, persist messages to an external store.
+const histories = new Map<string, ModelMessage[]>();
+
+function getHistory(sessionId: string): ModelMessage[] {
+  const existing = histories.get(sessionId);
+  if (existing) {
+    histories.delete(sessionId);
+    histories.set(sessionId, existing);
+    return existing;
+  }
+  if (histories.size >= HISTORY_LIMIT) {
+    const oldest = histories.keys().next().value;
+    if (oldest !== undefined) histories.delete(oldest);
+  }
+  const fresh: ModelMessage[] = [];
+  histories.set(sessionId, fresh);
+  return fresh;
+}
+
 const app = new BedrockAgentCoreApp({
   invocationHandler: {
     async *process(payload: any, context: any) {
+      const sessionId = context?.sessionId ?? 'default-session';
+      const history = getHistory(sessionId);
+      const userMessage: ModelMessage = { role: 'user', content: payload.prompt ?? '' };
+
       const model = await loadModel();
       const result = streamText({
         model,
         system: SYSTEM_PROMPT,
-        prompt: payload.prompt ?? '',
+        messages: [...history, userMessage],
       });
 
+      let assistant = '';
       for await (const chunk of result.textStream) {
+        assistant += chunk;
         yield { data: chunk };
       }
+
+      // Commit the exchange to history only after a non-empty reply. On a failed
+      // or empty stream the turn is dropped instead of leaving a dangling user
+      // (or empty assistant) message — consecutive same-role or empty-content
+      // messages would otherwise be rejected on the next turn for this session.
+      if (assistant.length > 0) {
+        history.push(userMessage, { role: 'assistant', content: assistant });
+      }
     },
   },
 });