aws-samples / aws-samples/sample-multi-agent-orchestration-chat-on-agentcore

feat: Speculative AgentCore Runtime pre-warm to reduce chat cold-start TTFB (fire-and-forget warmup invoke)

Open
#68 1 comment 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
TypeScript
Stars
128
Forks
12
Avg merge
3d 1h
Merged PRs (30d)
4

Description

## Summary

Reduce chat **TTFB (time to first token)** on cold sessions by adding a **speculative pre-warm** of the AgentCore Runtime microVM. When the user shows intent to chat (focuses / starts typing in the message input), the frontend fires a lightweight, fire-and-forget `warmup` invoke on the **same `sessionId`** so the platform provisions/assigns a microVM ahead of the first real message. The backend recognizes the `warmup` flag and **short-circuits** before any heavy work.

This adapts the "AgentCore runtime の warmup" technique (工夫 7) from the AWS Specialist Agent article to Moca's architecture:
- https://zenn.dev/aws_japan/articles/006-aws-specialist-agent

> ⚠️ Moca is a cost-optimized, multi-tenant deployment, so this proposal deliberately **does not** copy the article's `idleRuntimeSessionTimeout` extension (900s → 3600s). See "Moca-specific constraints" below.

## Background / Motivation

- The chat path talks to the AgentCore Runtime **directly** (not via the backend API):
- `packages/frontend/src/api/agent.ts` → `streamAgentResponse()` sets `X-Amzn-Bedrock-AgentCore-Runtime-Session-Id` and calls `agentClient.invoke()`.
- `packages/frontend/src/api/client/agent-client.ts` → `AgentClient.invoke()` posts to `VITE_AGENT_ENDPOINT` (the Runtime `/invocations` URL).
- `packages/frontend/src/api/client/base-client.ts` → `fetchWithAuth()` already attaches `Authorization: Bearer ` and `X-Amzn-Bedrock-AgentCore-Runtime-Custom-Id-Token: `.
- AgentCore routes invokes by session id (**microVM stickiness**). A brand-new `sessionId` pays the full startup path on its first invoke: microVM provision → ECR image pull → Node/Express boot → Strands SDK + tool/MCP registration. Sending a cheap `warmup` invoke first, with the **same** session-id header, lets the subsequent real invoke land on an already-warm microVM.
- `sessionId` is known *before* the first message in the common path: `NewChatRedirect` (`packages/frontend/src/components/NewChatRedirect.tsx`) generates it and navigates to `/chat/{sessionId}`; `useSessionSync` (`packages/frontend/src/hooks/useSessionSync.ts`) exposes it as `currentSessionId`. This gives us a natural window (user reading the welcome screen / typing) to warm in the background.

## Moca-specific constraints (differs from the article)

1. **Network mode = PUBLIC (no VPC).** The Runtime construct (`packages/cdk/lib/constructs/agentcore/agentcore-runtime.ts`) sets no VPC config, so cold start has **no ENI provisioning**. Warmup still helps (image pull + process boot + SDK/tool init), but the expected win is **smaller** than the article's VPC case. Set expectations accordingly and measure.
2. **Short idle timeout is intentional.** `idleRuntimeSessionTimeout` defaults to **300s** (`agentcore-runtime.ts`), and event/trigger sessions self-terminate (`packages/agent/src/services/session-terminator.ts`) — this is a deliberate cost posture after prior GB-hours/container-retention issues. **Do NOT raise it to 3600s** like the demo booth in the article.
- Implication: warm **close to send time** (on focus / first keystroke), not merely on page mount — otherwise a user who reads for > 5 min lets the warm microVM get reclaimed, wasting the warmup.
3. **Cost/UX tradeoff.** Every warmup provisions a microVM that bills memory until idle-reclaimed. A user who opens a chat and never sends = wasted spend. Mitigate with: intent-based trigger, **dedupe per session**, keep idle short, and a **feature flag** to disable.
4. **Reuse existing auth.** Unlike the article's hand-rolled `fetch`, Moca should route warmup through the existing `agentClient` / `fetchWithAuth` so the Bearer access token + id-token header are attached centrally (incl. the 401 refresh path).

## Requirements (EARS)

- The system shall provide a frontend `warmup(sessionId)` operation that issues a POST to the AgentCore Runtime `/invocations` endpoint carrying the `X-Amzn-Bedrock-AgentCore-Runtime-Session-Id: ` header and a body of `{ "warmup": true }`.
- The frontend shall trigger warmup on user intent (message-input **focus** and/or **first keystroke**) for the active `currentSessionId`.
- The frontend shall issue **at most one** warmup per `sessionId` (dedupe), and shall treat warmup as **fire-and-forget** (cancel the response body; never surface errors to the UI).
- When the backend receives an invocation whose body has `warmup === true`, the system shall respond `200 { "status": "warm" }` and shall **not** run prompt validation, identity resolution, session setup (AgentCore Memory / DynamoDB), agent construction, or model invocation.
- The backend shall only accept warmup from **authenticated** callers (i.e. the warmup short-circuit runs **after** JWT verification), so warmup cannot be used to spam microVM provisioning.
- The system shall not write any session history, metrics, or memory records for a warmup invoke.
- Where a warmup feature flag is disabled, the frontend shall not send any warmup requests.
- The system shall keep `idleRuntimeSessionTimeout` at its current cost-optimized value (300s) and shall not extend it as part of this change.

## Proposed approach (non-binding)

### Backend — `packages/agent`
- Add `warmup?: boolean` to `InvocationRequest` (`packages/agent/src/types/invocation-types.ts`).
- Add a `warmupMiddleware` and insert it in `packages/agent/src/app.ts` **after** `requestContextMiddleware` (JWT verified) and **before** `validateInvocationMiddleware` (which otherwise 400s on the empty prompt):
```
trackInFlight → requestContext → [warmup short-circuit] → validateInvocation → authResolver → identityResolver → handleInvocation
```
On `req.body.warmup === true`, respond `res.json({ status: 'warm' })` and return (do not call `next()`).
- Decide `trackInFlight` interaction: prefer that warmup does **not** hold the container "busy" beyond its instant response (it already releases on `finish`/`close`). Document the choice; a warmup invoke must never flip `/ping` to `HealthyBusy` for a meaningful window (`packages/agent/src/handlers/health.ts`).

### Frontend — `packages/frontend`
- Add `AgentClient.warmup(sessionId)` in `packages/frontend/src/api/client/agent-client.ts` (or a `warmup()` in `api/agent.ts`) that reuses `fetchWithAuth`, sets the session-id header, sends `{ warmup: true }`, then `await response.body?.cancel()` and swallows errors.
- Add a small `useRuntimeWarmup` hook (or wire directly into `packages/frontend/src/components/MessageInput.tsx`) that calls `warmup(currentSessionId)` on textarea **focus** / first `handleChange`, with a per-session dedupe `Set`/ref.
- Gate behind an env flag, e.g. `VITE_ENABLE_RUNTIME_WARMUP` (default on/off TBD).
- **Edge case:** a brand-new `/chat` with no id creates the id lazily in `onCreateSession()` on first send (`MessageInput.handleSubmit`). Since `NewChatRedirect` already routes to `/chat/{sessionId}` for the common entry paths, warm on focus for `currentSessionId`; document that the very first turn of a no-id `/chat` may not be pre-warmed (or pre-generate the id on focus and reuse it).

### CDK / config
- No change to `idleRuntimeSessionTimeout` (stays 300s). Optionally document the warmup flag in `docs/`.

## Acceptance criteria

- [ ] Frontend exposes a `warmup(sessionId)` that posts `{ warmup: true }` with the correct session-id header via the authenticated client, and cancels the response body.
- [ ] Warmup is triggered on message-input focus / first keystroke and is deduped to at most once per `sessionId`.
- [ ] Backend short-circuits `warmup === true` with `200 { status: 'warm' }` after JWT verification and before validation/identity/session/agent/model work.
- [ ] No session history / memory / metrics are written for a warmup invoke.
- [ ] A feature flag can fully disable warmup (no requests sent).
- [ ] `idleRuntimeSessionTimeout` remains 300s (unchanged).
- [ ] Unit tests: `warmupMiddleware` (short-circuit + auth-gated) and `agentClient.warmup` (headers/body/body-cancel + dedupe).
- [ ] A quick TTFB before/after measurement on a cold session is captured in the PR, and AgentCore Runtime GB-hours are sanity-checked for bounded cost impact.

## Out of scope

- Extending `idleRuntimeSessionTimeout` / changing the session self-termination policy.
- Moving to VPC network mode or changing the deploy artifact (code vs container).
- Pre-warming per-user identity credentials, MCP connections, or model clients (warmup stays lightweight, mirroring the article).

## References

- AWS Specialist Agent — 工夫 7「AgentCore runtime の warmup」: https://zenn.dev/aws_japan/articles/006-aws-specialist-agent
- AgentCore Runtime pre-warmed instances / speculative warmup (re:Post): https://repost.aws/articles/ARCJIn3t7aRC2FxiRTV1SuCA
- Relevant Moca files:
- Frontend: `packages/frontend/src/api/agent.ts`, `packages/frontend/src/api/client/agent-client.ts`, `packages/frontend/src/api/client/base-client.ts`, `packages/frontend/src/components/MessageInput.tsx`, `packages/frontend/src/hooks/useSessionSync.ts`, `packages/frontend/src/components/NewChatRedirect.tsx`
- Backend: `packages/agent/src/app.ts`, `packages/agent/src/handlers/invocations.ts`, `packages/agent/src/handlers/health.ts`, `packages/agent/src/libs/middleware/{request-context,validate-invocation,track-in-flight}.ts`, `packages/agent/src/types/invocation-types.ts`
- CDK: `packages/cdk/lib/constructs/agentcore/agentcore-runtime.ts` (`idleRuntimeSessionTimeout`)

Contributor guide

Open the contributing guide

Research direction

Start with the middleware order in packages/agent/src/app.ts and the request types in packages/agent/src/types/invocation-types.ts, then inspect AgentClient and MessageInput in packages/frontend. Add the authenticated warmup path, intent trigger, dedupe, feature flag, and middleware tests described in the acceptance criteria. Done means warmup is auth-gated and side-effect-free, the response body is cancelled, idleRuntimeSessionTimeout remains 300s, and cold-session TTFB and GB-hours are measured.

Written by the indexing model from the issue text.

Assessment

Tech stack
aws, express, typescript
Domain
api, backend, frontend, performance, testing
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
50/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.