From 3041029766c5b4df199e6abc87de1e776acb9dfc Mon Sep 17 00:00:00 2001 From: jariy17 Date: Mon, 3 Aug 2026 20:10:55 +0000 Subject: [PATCH] fix: validate generated agent inputs --- docs/frameworks.md | 5 + e2e-tests/fixtures/import/app/main.py | 5 +- src/assets/README.md | 5 + .../assets.snapshot.test.ts.snap | 118 ++++++++++++++++-- src/assets/__tests__/input-validation.test.ts | 44 +++++++ src/assets/agents/AGENTS.md | 2 + src/assets/python/http/autogen/base/README.md | 5 + src/assets/python/http/autogen/base/main.py | 2 + .../python/http/googleadk/base/README.md | 5 + src/assets/python/http/googleadk/base/main.py | 2 + .../http/langchain_langgraph/base/README.md | 5 + .../http/langchain_langgraph/base/main.py | 4 + .../python/http/openaiagents/base/README.md | 5 + .../python/http/openaiagents/base/main.py | 2 + src/assets/python/http/strands/base/README.md | 6 + src/assets/python/http/strands/base/main.py | 44 ++++++- .../typescript/http/strands/base/README.md | 5 + .../typescript/http/strands/base/main.ts | 11 +- .../typescript/http/vercelai/base/README.md | 5 + .../typescript/http/vercelai/base/main.ts | 10 +- .../agent/import/__tests__/translator.test.ts | 2 + .../agent/import/base-translator.ts | 2 + .../__tests__/wirePaymentCapability.test.ts | 10 +- 23 files changed, 283 insertions(+), 21 deletions(-) create mode 100644 src/assets/__tests__/input-validation.test.ts diff --git a/docs/frameworks.md b/docs/frameworks.md index 67ce4f32c..d3bb2206b 100644 --- a/docs/frameworks.md +++ b/docs/frameworks.md @@ -24,6 +24,11 @@ framework is restricted to `Strands` or `VercelAI`; other values are rejected. S | **OpenAIAgents** | OpenAI only | | **VercelAI** | Bedrock, Anthropic, OpenAI, Gemini | +## Runtime Input Validation + +Validate agent invocation payloads before passing them to a framework. Template entrypoints validate plain prompts as +strings; preserve that validation when extending a generated application, and pass only prompt text to the agent. + ## Framework Selection Guide ### Strands Agents diff --git a/e2e-tests/fixtures/import/app/main.py b/e2e-tests/fixtures/import/app/main.py index 2400ee741..f5fc65658 100644 --- a/e2e-tests/fixtures/import/app/main.py +++ b/e2e-tests/fixtures/import/app/main.py @@ -34,7 +34,10 @@ def get_or_create_agent(): async def invoke(payload, context): log.info("Invoking Agent.....") agent = get_or_create_agent() - stream = agent.stream_async(payload.get("prompt")) + prompt = payload.get("prompt") + if not isinstance(prompt, str): + raise ValueError("prompt must be a string") + stream = agent.stream_async(prompt) async for event in stream: if "data" in event and isinstance(event["data"], str): yield event["data"] diff --git a/src/assets/README.md b/src/assets/README.md index 7064d6466..b3ed48ce2 100644 --- a/src/assets/README.md +++ b/src/assets/README.md @@ -37,6 +37,11 @@ Run your agent locally: agentcore dev ``` +### Validate Invocation Input + +Validate runtime invocation payloads before forwarding them to an agent framework. Keep user prompts typed as strings +and pass only prompt text to the agent. + ### Deployment Deploy to AWS: diff --git a/src/assets/__tests__/__snapshots__/assets.snapshot.test.ts.snap b/src/assets/__tests__/__snapshots__/assets.snapshot.test.ts.snap index 67e1d31d2..ca516e32d 100644 --- a/src/assets/__tests__/__snapshots__/assets.snapshot.test.ts.snap +++ b/src/assets/__tests__/__snapshots__/assets.snapshot.test.ts.snap @@ -3335,6 +3335,11 @@ file defines a Starlette ASGI app with the AutoGen framework running within. \`model/load.py\` instantiates your chosen model provider. +## Input Validation + +Validate invocation input before forwarding it to the agent framework. Keep plain prompts typed as strings and pass +only prompt text to the agent. + ## Environment Variables | Variable | Required | Description | @@ -3538,6 +3543,8 @@ async def invoke(payload, context): # Process the user prompt prompt = payload.get("prompt", "What can you help me with?") + if not isinstance(prompt, str): + raise ValueError("prompt must be a string") session_id = getattr(context, "session_id", "default-session") # Reuse the per-session agent (preserves conversation history) @@ -3791,6 +3798,11 @@ file defines a Starlette ASGI app with the Google ADK framework running within. \`model/load.py\` instantiates your chosen model provider (Gemini). +## Input Validation + +Validate invocation input before forwarding it to the agent framework. Keep plain prompts typed as strings and pass +only prompt text to the agent. + ## Environment Variables | Variable | Required | Description | @@ -4058,6 +4070,8 @@ async def invoke(payload, context): # Process the user prompt prompt = payload.get("prompt", "What can you help me with?") + if not isinstance(prompt, str): + raise ValueError("prompt must be a string") session_id = getattr(context, "session_id", "default_session") user_id = payload.get("user_id", "default_user") @@ -4251,6 +4265,11 @@ file defines a Starlette ASGI app with the LangChain/LangGraph framework running \`model/load.py\` instantiates your chosen model provider. +## Input Validation + +Validate invocation input before forwarding it to the agent framework. Keep plain prompts typed as strings and pass +only prompt text to the agent. + ## Environment Variables | Variable | Required | Description | @@ -4510,6 +4529,8 @@ async def invoke(payload, context): # Process the user prompt prompt = payload.get("prompt", "What can you help me with?") + if not isinstance(prompt, str): + raise ValueError("prompt must be a string") session_id = getattr(context, "session_id", "default-session") touch_thread(session_id) log.info(f"Agent input: {prompt}") @@ -4529,6 +4550,8 @@ async def invoke(payload, context): # Process the user prompt prompt = payload.get("prompt", "What can you help me with?") + if not isinstance(prompt, str): + raise ValueError("prompt must be a string") session_id = getattr(context, "session_id", "default-session") touch_thread(session_id) log.info(f"Agent input: {prompt}") @@ -4829,6 +4852,11 @@ file defines a Starlette ASGI app with the OpenAI Agents SDK framework running w \`model/load.py\` instantiates your chosen model provider (OpenAI). +## Input Validation + +Validate invocation input before forwarding it to the agent framework. Keep plain prompts typed as strings and pass +only prompt text to the agent. + ## Environment Variables | Variable | Required | Description | @@ -5089,6 +5117,8 @@ async def invoke(payload, context): # Process the user prompt prompt = payload.get("prompt", "What can you help me with?") + if not isinstance(prompt, str): + raise ValueError("prompt must be a string") session_id = getattr(context, "session_id", "default-session") session = get_session(session_id) @@ -5274,6 +5304,12 @@ file defines a Starlette ASGI app with the chosen Agent framework SDK running wi \`model/load.py\` instantiates your chosen model provider. +## Input Validation + +Validate invocation input before forwarding it to Strands. Keep plain prompts typed as strings. If the app accepts a +caller-supplied message history, retain \`strip_trailing_tool_use()\`, which normalizes the history tail before +invoking the agent. + ## Environment Variables | Variable | Required | Description | @@ -5830,17 +5866,53 @@ get_or_create_agent = agent_factory() {{/if}} +def strip_trailing_tool_use(messages: Any) -> list[dict]: + """Strip toolUse blocks from the tail until the last message has none.""" + if not isinstance(messages, list): + raise ValueError("messages must be a list") + + messages = list(messages) + while messages: + last = messages[-1] + if not isinstance(last, dict): + raise ValueError("each message must be an object") + original_content = last.get("content", []) + if not isinstance(original_content, list) or not all(isinstance(block, dict) for block in original_content): + raise ValueError("each message content value must be a list of content blocks") + + content = [block for block in original_content if "toolUse" not in block] + if len(content) == len(original_content): + break + if content: + messages[-1] = {**last, "content": content} + break + messages.pop() + + return messages + + def _extract_prompt(payload: dict): - """Accept harness-style messages[], tool_results[], or plain prompt string payloads.""" + """Accept validated harness messages, tool results, or a plain prompt string.""" + if not isinstance(payload, dict): + raise ValueError("payload must be a JSON object") if "messages" in payload: - return payload["messages"] + return strip_trailing_tool_use(payload["messages"]) if "tool_results" in payload: + tool_results = payload["tool_results"] + if not isinstance(tool_results, list) or not all( + isinstance(tool_result, dict) and isinstance(tool_result.get("toolUseId"), str) + for tool_result in tool_results + ): + raise ValueError("tool_results must contain objects with a toolUseId string") return [{"role": "user", "content": [{"toolResult": { "toolUseId": tr["toolUseId"], "status": tr.get("status", "success"), "content": tr.get("content", []), - }} for tr in payload["tool_results"]]}] - return payload.get("prompt", "") + }} for tr in tool_results]}] + prompt = payload.get("prompt", "") + if not isinstance(prompt, str): + raise ValueError("prompt must be a string") + return prompt def _has_inline_function_call(messages) -> bool: @@ -7308,6 +7380,11 @@ Run your agent locally: agentcore dev \`\`\` +### Validate Invocation Input + +Validate runtime invocation payloads before forwarding them to an agent framework. Keep user prompts typed as strings +and pass only prompt text to the agent. + ### Deployment Deploy to AWS: @@ -7402,6 +7479,8 @@ Tags defined in \`agentcore.json\` flow through to deployed CloudFormation resou \`agentcore validate\` to check. 4. **Resource Removal:** Use \`agentcore remove\` to remove resources. Run \`agentcore deploy\` after removal to tear down deployed infrastructure. +5. **Invocation Input:** Validate runtime payloads and require text prompts to be strings. If a Strands app accepts a + caller-supplied message history, normalize the history tail with \`strip_trailing_tool_use()\` before invocation. ## Directory Structure @@ -7615,6 +7694,11 @@ defines an HTTP server that streams tokens from your chosen Agent framework SDK. \`model/load.ts\` instantiates your chosen model provider. +## Input Validation + +The generated Zod request schema keeps plain prompts typed as strings before forwarding them to Strands. Retain this +validation when extending the request shape, and pass only prompt text to the agent. + ## Environment Variables | Variable | Required | Description | @@ -7703,6 +7787,10 @@ const SYSTEM_PROMPT = \` You are a helpful assistant. Use tools when appropriate. \`; +const requestSchema = z.object({ + prompt: z.string().default(''), +}); + {{#if hasMemory}} const agentCache = new Map(); @@ -7756,7 +7844,8 @@ async function getOrCreateAgent(sessionId: string): Promise { const app = new BedrockAgentCoreApp({ invocationHandler: { - async *process(payload: any, context: any) { + requestSchema, + async *process(payload, context) { {{#if hasMemory}} const sessionId = context?.sessionId ?? 'default-session'; const actorId = getActorId(payload, context); @@ -7768,7 +7857,7 @@ const app = new BedrockAgentCoreApp({ {{#if hasMemory}} try { - for await (const event of agent.stream(payload.prompt ?? '')) { + for await (const event of agent.stream(payload.prompt)) { if ( event.type === 'modelStreamUpdateEvent' && event.event?.type === 'modelContentBlockDeltaEvent' && @@ -7792,7 +7881,7 @@ const app = new BedrockAgentCoreApp({ // e.g. Anthropic). Restoring on error keeps the session reusable. const snapshot = agent.takeSnapshot({ include: ['messages'] }); try { - for await (const event of agent.stream(payload.prompt ?? '')) { + for await (const event of agent.stream(payload.prompt)) { if ( event.type === 'modelStreamUpdateEvent' && event.event?.type === 'modelContentBlockDeltaEvent' && @@ -8072,6 +8161,11 @@ defines an HTTP app that streams tokens using the Vercel AI SDK's \`streamText\` \`model/load.ts\` instantiates your chosen model provider. +## Input Validation + +The generated Zod request schema keeps plain prompts typed as strings before forwarding them to the agent framework. +Retain this validation when extending the request shape, and pass only prompt text to the agent. + ## Environment Variables | Variable | Required | Description | @@ -8126,10 +8220,15 @@ Thumbs.db exports[`Assets Directory Snapshots > TypeScript assets > typescript/typescript/http/vercelai/base/main.ts should match snapshot 1`] = ` "import { BedrockAgentCoreApp } from 'bedrock-agentcore/runtime'; import { streamText, type ModelMessage } from 'ai'; +import { z } from 'zod'; import { loadModel } from './model/load.js'; const SYSTEM_PROMPT = \`You are a helpful assistant.\`; +const requestSchema = z.object({ + prompt: z.string().default(''), +}); + const HISTORY_LIMIT = 128; // Keeps one message history per sessionId so each session remembers its own @@ -8158,10 +8257,11 @@ function getHistory(sessionId: string): ModelMessage[] { const app = new BedrockAgentCoreApp({ invocationHandler: { - async *process(payload: any, context: any) { + requestSchema, + async *process(payload, context) { const sessionId = context?.sessionId ?? 'default-session'; const history = getHistory(sessionId); - const userMessage: ModelMessage = { role: 'user', content: payload.prompt ?? '' }; + const userMessage: ModelMessage = { role: 'user', content: payload.prompt }; const model = await loadModel(); const result = streamText({ diff --git a/src/assets/__tests__/input-validation.test.ts b/src/assets/__tests__/input-validation.test.ts new file mode 100644 index 000000000..2ae10d3aa --- /dev/null +++ b/src/assets/__tests__/input-validation.test.ts @@ -0,0 +1,44 @@ +import { readFileSync } from 'node:fs'; +import { resolve } from 'node:path'; +import { describe, expect, it } from 'vitest'; + +const ASSETS_DIR = resolve(__dirname, '..'); + +const PYTHON_HTTP_ENTRYPOINTS = [ + 'python/http/autogen/base/main.py', + 'python/http/googleadk/base/main.py', + 'python/http/langchain_langgraph/base/main.py', + 'python/http/openaiagents/base/main.py', + 'python/http/strands/base/main.py', +]; + +const TYPESCRIPT_HTTP_ENTRYPOINTS = [ + 'typescript/http/strands/base/main.ts', + 'typescript/http/vercelai/base/main.ts', +]; + +describe('HTTP agent template input validation', () => { + it.each(PYTHON_HTTP_ENTRYPOINTS)('%s rejects non-string prompts', templatePath => { + const template = readFileSync(resolve(ASSETS_DIR, templatePath), 'utf8'); + + expect(template).toContain('if not isinstance(prompt, str):'); + expect(template).toContain('raise ValueError("prompt must be a string")'); + }); + + it.each(TYPESCRIPT_HTTP_ENTRYPOINTS)('%s validates prompts with Zod', templatePath => { + const template = readFileSync(resolve(ASSETS_DIR, templatePath), 'utf8'); + + expect(template).toContain("prompt: z.string().default('')"); + expect(template).toContain('requestSchema,'); + }); + + it('strips toolUse blocks from the Python Strands message-history tail', () => { + const template = readFileSync(resolve(ASSETS_DIR, 'python/http/strands/base/main.py'), 'utf8'); + + expect(template).toContain('def strip_trailing_tool_use(messages: Any) -> list[dict]:'); + expect(template).toContain('while messages:'); + expect(template).toContain('content = [block for block in original_content if "toolUse" not in block]'); + expect(template).toContain('messages.pop()'); + expect(template).toContain('return strip_trailing_tool_use(payload["messages"])'); + }); +}); diff --git a/src/assets/agents/AGENTS.md b/src/assets/agents/AGENTS.md index 53877021e..10c39e285 100644 --- a/src/assets/agents/AGENTS.md +++ b/src/assets/agents/AGENTS.md @@ -23,6 +23,8 @@ Tags defined in `agentcore.json` flow through to deployed CloudFormation resourc `agentcore validate` to check. 4. **Resource Removal:** Use `agentcore remove` to remove resources. Run `agentcore deploy` after removal to tear down deployed infrastructure. +5. **Invocation Input:** Validate runtime payloads and require text prompts to be strings. If a Strands app accepts a + caller-supplied message history, normalize the history tail with `strip_trailing_tool_use()` before invocation. ## Directory Structure diff --git a/src/assets/python/http/autogen/base/README.md b/src/assets/python/http/autogen/base/README.md index 124380124..1172df7ce 100644 --- a/src/assets/python/http/autogen/base/README.md +++ b/src/assets/python/http/autogen/base/README.md @@ -13,6 +13,11 @@ file defines a Starlette ASGI app with the AutoGen framework running within. `model/load.py` instantiates your chosen model provider. +## Input Validation + +Validate invocation input before forwarding it to the agent framework. Keep plain prompts typed as strings and pass +only prompt text to the agent. + ## Environment Variables | Variable | Required | Description | diff --git a/src/assets/python/http/autogen/base/main.py b/src/assets/python/http/autogen/base/main.py index 4c0beb3b4..a6d45a47b 100644 --- a/src/assets/python/http/autogen/base/main.py +++ b/src/assets/python/http/autogen/base/main.py @@ -127,6 +127,8 @@ async def invoke(payload, context): # Process the user prompt prompt = payload.get("prompt", "What can you help me with?") + if not isinstance(prompt, str): + raise ValueError("prompt must be a string") session_id = getattr(context, "session_id", "default-session") # Reuse the per-session agent (preserves conversation history) diff --git a/src/assets/python/http/googleadk/base/README.md b/src/assets/python/http/googleadk/base/README.md index 93c760820..65a3f0532 100644 --- a/src/assets/python/http/googleadk/base/README.md +++ b/src/assets/python/http/googleadk/base/README.md @@ -13,6 +13,11 @@ file defines a Starlette ASGI app with the Google ADK framework running within. `model/load.py` instantiates your chosen model provider (Gemini). +## Input Validation + +Validate invocation input before forwarding it to the agent framework. Keep plain prompts typed as strings and pass +only prompt text to the agent. + ## Environment Variables | Variable | Required | Description | diff --git a/src/assets/python/http/googleadk/base/main.py b/src/assets/python/http/googleadk/base/main.py index 165ab7e96..73abc4323 100644 --- a/src/assets/python/http/googleadk/base/main.py +++ b/src/assets/python/http/googleadk/base/main.py @@ -191,6 +191,8 @@ async def invoke(payload, context): # Process the user prompt prompt = payload.get("prompt", "What can you help me with?") + if not isinstance(prompt, str): + raise ValueError("prompt must be a string") session_id = getattr(context, "session_id", "default_session") user_id = payload.get("user_id", "default_user") diff --git a/src/assets/python/http/langchain_langgraph/base/README.md b/src/assets/python/http/langchain_langgraph/base/README.md index 336477891..9787803db 100644 --- a/src/assets/python/http/langchain_langgraph/base/README.md +++ b/src/assets/python/http/langchain_langgraph/base/README.md @@ -13,6 +13,11 @@ file defines a Starlette ASGI app with the LangChain/LangGraph framework running `model/load.py` instantiates your chosen model provider. +## Input Validation + +Validate invocation input before forwarding it to the agent framework. Keep plain prompts typed as strings and pass +only prompt text to the agent. + ## Environment Variables | Variable | Required | Description | diff --git a/src/assets/python/http/langchain_langgraph/base/main.py b/src/assets/python/http/langchain_langgraph/base/main.py index 387fe2870..f70ee9b7b 100644 --- a/src/assets/python/http/langchain_langgraph/base/main.py +++ b/src/assets/python/http/langchain_langgraph/base/main.py @@ -183,6 +183,8 @@ async def invoke(payload, context): # Process the user prompt prompt = payload.get("prompt", "What can you help me with?") + if not isinstance(prompt, str): + raise ValueError("prompt must be a string") session_id = getattr(context, "session_id", "default-session") touch_thread(session_id) log.info(f"Agent input: {prompt}") @@ -202,6 +204,8 @@ async def invoke(payload, context): # Process the user prompt prompt = payload.get("prompt", "What can you help me with?") + if not isinstance(prompt, str): + raise ValueError("prompt must be a string") session_id = getattr(context, "session_id", "default-session") touch_thread(session_id) log.info(f"Agent input: {prompt}") diff --git a/src/assets/python/http/openaiagents/base/README.md b/src/assets/python/http/openaiagents/base/README.md index ea9b2eb12..2258553d6 100644 --- a/src/assets/python/http/openaiagents/base/README.md +++ b/src/assets/python/http/openaiagents/base/README.md @@ -13,6 +13,11 @@ file defines a Starlette ASGI app with the OpenAI Agents SDK framework running w `model/load.py` instantiates your chosen model provider (OpenAI). +## Input Validation + +Validate invocation input before forwarding it to the agent framework. Keep plain prompts typed as strings and pass +only prompt text to the agent. + ## Environment Variables | Variable | Required | Description | diff --git a/src/assets/python/http/openaiagents/base/main.py b/src/assets/python/http/openaiagents/base/main.py index db1a43d00..868144a8a 100644 --- a/src/assets/python/http/openaiagents/base/main.py +++ b/src/assets/python/http/openaiagents/base/main.py @@ -184,6 +184,8 @@ async def invoke(payload, context): # Process the user prompt prompt = payload.get("prompt", "What can you help me with?") + if not isinstance(prompt, str): + raise ValueError("prompt must be a string") session_id = getattr(context, "session_id", "default-session") session = get_session(session_id) diff --git a/src/assets/python/http/strands/base/README.md b/src/assets/python/http/strands/base/README.md index cae68f70c..1985f1eae 100644 --- a/src/assets/python/http/strands/base/README.md +++ b/src/assets/python/http/strands/base/README.md @@ -13,6 +13,12 @@ file defines a Starlette ASGI app with the chosen Agent framework SDK running wi `model/load.py` instantiates your chosen model provider. +## Input Validation + +Validate invocation input before forwarding it to Strands. Keep plain prompts typed as strings. If the app accepts a +caller-supplied message history, retain `strip_trailing_tool_use()`, which normalizes the history tail before +invoking the agent. + ## Environment Variables | Variable | Required | Description | diff --git a/src/assets/python/http/strands/base/main.py b/src/assets/python/http/strands/base/main.py index 5676b86d9..69dca7e8e 100644 --- a/src/assets/python/http/strands/base/main.py +++ b/src/assets/python/http/strands/base/main.py @@ -481,17 +481,53 @@ def get_or_create_agent(session_id{{#if hasSkillsFetcher}}, skill_plugins=None{{ {{/if}} +def strip_trailing_tool_use(messages: Any) -> list[dict]: + """Strip toolUse blocks from the tail until the last message has none.""" + if not isinstance(messages, list): + raise ValueError("messages must be a list") + + messages = list(messages) + while messages: + last = messages[-1] + if not isinstance(last, dict): + raise ValueError("each message must be an object") + original_content = last.get("content", []) + if not isinstance(original_content, list) or not all(isinstance(block, dict) for block in original_content): + raise ValueError("each message content value must be a list of content blocks") + + content = [block for block in original_content if "toolUse" not in block] + if len(content) == len(original_content): + break + if content: + messages[-1] = {**last, "content": content} + break + messages.pop() + + return messages + + def _extract_prompt(payload: dict): - """Accept harness-style messages[], tool_results[], or plain prompt string payloads.""" + """Accept validated harness messages, tool results, or a plain prompt string.""" + if not isinstance(payload, dict): + raise ValueError("payload must be a JSON object") if "messages" in payload: - return payload["messages"] + return strip_trailing_tool_use(payload["messages"]) if "tool_results" in payload: + tool_results = payload["tool_results"] + if not isinstance(tool_results, list) or not all( + isinstance(tool_result, dict) and isinstance(tool_result.get("toolUseId"), str) + for tool_result in tool_results + ): + raise ValueError("tool_results must contain objects with a toolUseId string") return [{"role": "user", "content": [{"toolResult": { "toolUseId": tr["toolUseId"], "status": tr.get("status", "success"), "content": tr.get("content", []), - }} for tr in payload["tool_results"]]}] - return payload.get("prompt", "") + }} for tr in tool_results]}] + prompt = payload.get("prompt", "") + if not isinstance(prompt, str): + raise ValueError("prompt must be a string") + return prompt def _has_inline_function_call(messages) -> bool: diff --git a/src/assets/typescript/http/strands/base/README.md b/src/assets/typescript/http/strands/base/README.md index 69903b4ac..9f7e68959 100644 --- a/src/assets/typescript/http/strands/base/README.md +++ b/src/assets/typescript/http/strands/base/README.md @@ -13,6 +13,11 @@ defines an HTTP server that streams tokens from your chosen Agent framework SDK. `model/load.ts` instantiates your chosen model provider. +## Input Validation + +The generated Zod request schema keeps plain prompts typed as strings before forwarding them to Strands. Retain this +validation when extending the request shape, and pass only prompt text to the agent. + ## Environment Variables | Variable | Required | Description | diff --git a/src/assets/typescript/http/strands/base/main.ts b/src/assets/typescript/http/strands/base/main.ts index fe7c7136b..9722285e5 100644 --- a/src/assets/typescript/http/strands/base/main.ts +++ b/src/assets/typescript/http/strands/base/main.ts @@ -34,6 +34,10 @@ const SYSTEM_PROMPT = ` You are a helpful assistant. Use tools when appropriate. `; +const requestSchema = z.object({ + prompt: z.string().default(''), +}); + {{#if hasMemory}} const agentCache = new Map(); @@ -87,7 +91,8 @@ async function getOrCreateAgent(sessionId: string): Promise { const app = new BedrockAgentCoreApp({ invocationHandler: { - async *process(payload: any, context: any) { + requestSchema, + async *process(payload, context) { {{#if hasMemory}} const sessionId = context?.sessionId ?? 'default-session'; const actorId = getActorId(payload, context); @@ -99,7 +104,7 @@ const app = new BedrockAgentCoreApp({ {{#if hasMemory}} try { - for await (const event of agent.stream(payload.prompt ?? '')) { + for await (const event of agent.stream(payload.prompt)) { if ( event.type === 'modelStreamUpdateEvent' && event.event?.type === 'modelContentBlockDeltaEvent' && @@ -123,7 +128,7 @@ const app = new BedrockAgentCoreApp({ // e.g. Anthropic). Restoring on error keeps the session reusable. const snapshot = agent.takeSnapshot({ include: ['messages'] }); try { - for await (const event of agent.stream(payload.prompt ?? '')) { + for await (const event of agent.stream(payload.prompt)) { if ( event.type === 'modelStreamUpdateEvent' && event.event?.type === 'modelContentBlockDeltaEvent' && diff --git a/src/assets/typescript/http/vercelai/base/README.md b/src/assets/typescript/http/vercelai/base/README.md index 7b1d8e0e3..899dc80c1 100644 --- a/src/assets/typescript/http/vercelai/base/README.md +++ b/src/assets/typescript/http/vercelai/base/README.md @@ -13,6 +13,11 @@ defines an HTTP app that streams tokens using the Vercel AI SDK's `streamText` A `model/load.ts` instantiates your chosen model provider. +## Input Validation + +The generated Zod request schema keeps plain prompts typed as strings before forwarding them to the agent framework. +Retain this validation when extending the request shape, and pass only prompt text to the agent. + ## Environment Variables | Variable | Required | Description | diff --git a/src/assets/typescript/http/vercelai/base/main.ts b/src/assets/typescript/http/vercelai/base/main.ts index b899c9f31..8051fea01 100644 --- a/src/assets/typescript/http/vercelai/base/main.ts +++ b/src/assets/typescript/http/vercelai/base/main.ts @@ -1,9 +1,14 @@ import { BedrockAgentCoreApp } from 'bedrock-agentcore/runtime'; import { streamText, type ModelMessage } from 'ai'; +import { z } from 'zod'; import { loadModel } from './model/load.js'; const SYSTEM_PROMPT = `You are a helpful assistant.`; +const requestSchema = z.object({ + prompt: z.string().default(''), +}); + const HISTORY_LIMIT = 128; // Keeps one message history per sessionId so each session remembers its own @@ -32,10 +37,11 @@ function getHistory(sessionId: string): ModelMessage[] { const app = new BedrockAgentCoreApp({ invocationHandler: { - async *process(payload: any, context: any) { + requestSchema, + async *process(payload, context) { const sessionId = context?.sessionId ?? 'default-session'; const history = getHistory(sessionId); - const userMessage: ModelMessage = { role: 'user', content: payload.prompt ?? '' }; + const userMessage: ModelMessage = { role: 'user', content: payload.prompt }; const model = await loadModel(); const result = streamText({ diff --git a/src/cli/operations/agent/import/__tests__/translator.test.ts b/src/cli/operations/agent/import/__tests__/translator.test.ts index 5090ebbd6..5e8b2ad1c 100644 --- a/src/cli/operations/agent/import/__tests__/translator.test.ts +++ b/src/cli/operations/agent/import/__tests__/translator.test.ts @@ -54,6 +54,8 @@ describe('StrandsTranslator', () => { expect(result.mainPyContent).toContain('def invoke_agent(question: str'); expect(result.mainPyContent).toContain('@app.entrypoint'); expect(result.mainPyContent).toContain('async def invoke(payload, context):'); + expect(result.mainPyContent).toContain('if not isinstance(agent_query, str):'); + expect(result.mainPyContent).toContain('raise ValueError("prompt must be a string")'); expect(result.collaboratorFiles.size).toBe(0); expect(result.features.hasMemory).toBe(false); expect(result.features.hasActionGroups).toBe(false); diff --git a/src/cli/operations/agent/import/base-translator.ts b/src/cli/operations/agent/import/base-translator.ts index daecb08f6..47fa7ffcb 100644 --- a/src/cli/operations/agent/import/base-translator.ts +++ b/src/cli/operations/agent/import/base-translator.ts @@ -326,6 +326,8 @@ memory_id = os.environ.get("MEMORY_ID", "") '', ' tools_used.clear()', ' agent_query = payload.get("prompt", "")', + ' if not isinstance(agent_query, str):', + ' raise ValueError("prompt must be a string")', ' if not agent_query:', ' yield "No query provided, please provide a \'prompt\' field in the payload."', ' return', diff --git a/src/cli/primitives/__tests__/wirePaymentCapability.test.ts b/src/cli/primitives/__tests__/wirePaymentCapability.test.ts index d00fbf2fa..1e641a638 100644 --- a/src/cli/primitives/__tests__/wirePaymentCapability.test.ts +++ b/src/cli/primitives/__tests__/wirePaymentCapability.test.ts @@ -166,7 +166,10 @@ describe('wirePaymentCapability (via PaymentManagerPrimitive.add)', () => { '@app.entrypoint', 'async def invoke(payload, context):', ' agent = get_or_create_agent()', - ' stream = agent.stream_async(payload.get("prompt"))', + ' prompt = payload.get("prompt")', + ' if not isinstance(prompt, str):', + ' raise ValueError("prompt must be a string")', + ' stream = agent.stream_async(prompt)', ' async for event in stream:', ' yield event', ].join('\n'); @@ -242,7 +245,10 @@ describe('wirePaymentCapability (via PaymentManagerPrimitive.add)', () => { ' system_prompt="You are a payment assistant.",', ' tools=my_tools,', ' )', - ' stream = agent.stream_async(payload.get("prompt"))', + ' prompt = payload.get("prompt")', + ' if not isinstance(prompt, str):', + ' raise ValueError("prompt must be a string")', + ' stream = agent.stream_async(prompt)', ' async for event in stream:', ' yield event', ].join('\n');