Lessons Learned: Shipping a Koog + Bedrock AI Agent to Production on AWS Lambda
Everything the tutorials don't tell you - from an AWS 10,000 AIdeas-winning project.
Series: Kotlin on AWS (1 article)
- 1Lessons Learned: Shipping a Koog + Bedrock AI Agent to Production on AWS Lambda This article
Lead Software Engineer | AWS Community Builder
I built MockNest Serverless — an open-source AI-powered mock server that generates WireMock mocks from API specifications. It runs on AWS Lambda, uses Koog [1] (JetBrains’ Kotlin AI agent framework) for the agent logic, and uses Amazon Bedrock with Nova Pro as the brain. It was selected as an AWS 10,000 AIdeas Finalist and later became one of the winning projects, which was pretty cool [15].
Along the way I learnt a bunch of things the hard way. This article is about those lessons.
Lesson 1: The Tutorial Setup Will Betray You in Production
The Koog Bedrock notebook example [2] uses
simpleBedrockExecutor with access keys from environment variables.That’s fine for a local Notebooks demo, but do not ship this to production. Although Lambda environment variables are encrypted at rest, users with sufficient permissions can still view them in plain text. And you don’t need credentials at all on Lambda — the execution role already provides temporary credentials that the SDK picks up automatically.
Use Koog’s API with the Kotlin AWS SDK’s
BedrockRuntimeClient instead — no credentials in sight:val bedrockClient = BedrockRuntimeClient { region = "eu-west-1" }val bedrockLLMClient = BedrockLLMClient(bedrockClient, apiMethod = BedrockAPIMethod.Converse)val executor = MultiLLMPromptExecutor(LLMProvider.Bedrock to bedrockLLMClient)Same Koog executor, same agent, same everything downstream.
BedrockAPIMethod.Converse uses the Converse API [3] — AWS's recommended unified interface across all Bedrock models. It handles model-specific prompt formatting and returns token usage metadata in the response, which is useful for cost tracking.
Your Lambda role needs Bedrock permissions. Here’s a SAM example (same applies for CDK, Terraform, etc.):- PolicyName: BedrockAccess PolicyDocument: Version: '2012-10-17' Statement: - Effect: Allow Action: - bedrock:InvokeModel - bedrock:InvokeModelWithResponseStream Resource: - !Sub "arn:aws:bedrock:${AWS::Region}::foundation-model/*" - !Sub "arn:aws:bedrock:${AWS::Region}::inference-profile/*"If you use cross-region inference profiles, you may need to widen the region and add access to foundation-model resources in target regions.
This also matters for SnapStart [5]. SnapStart restores Lambda from an initialization snapshot, and AWS says connections created during initialization are not guaranteed after restore. In my case, the Bedrock path was initialized too eagerly during priming, so connection-related state ended up in the snapshot unintentionally. During performance tests, when Lambda had to scale out and restore new environments, that showed up as restore-time timeouts. The fix was to keep the Bedrock path lazy: no explicit static credentials and no real Bedrock calls during priming. When you don’t set
credentialsProvider , the Kotlin AWS SDK uses the default credentials provider chain [4], so credentials are resolved when requests are actually made rather than being wired into client initialization.Lesson 2: The Region Puzzle
You write
BedrockModels.AmazonNovaPro and get an error. Depending on your region and setup, you may use the direct model ID or an inference profile. Inference profiles are especially useful for cross-region routing:val model = BedrockModels.AmazonNovaPro.withInferenceProfile("eu")// Produces: eu.amazon.nova-pro-v1:0Deploy to
eu-west-1 ? Use "eu" . Deploy to us-east-1 ? Use "us" . Full mapping is in the AWS docs [6].Watch out for model access too. Some Bedrock models require extra setup before you can use them. Some (like Amazon Nova) may already be available. Others require provider-specific setup — check the Bedrock console [7] under Model catalog. I got an
AccessDeniedException the first time I tried Anthropic Claude because I hadn't submitted the required use case form. Not an IAM issue — a licensing step that's easy to miss.Lesson 3: Building a Multi-Step Pipeline with Koog’s Strategy Graph
Most real-world LLM use cases aren’t a single call. You want a pipeline: call the LLM, validate the output, and if something’s off, send it back with feedback. Koog’s strategy DSL is a directed graph of nodes with conditional edges — it manages the LLM session across nodes and keeps the wiring separate from the logic.
val mockGenerationStrategy = strategy<SpecWithDescriptionRequest, GenerationResult>("mock-generation") {
val setupNode by node<SpecWithDescriptionRequest, MockGenerationContext>("setup") { request -> val specification = specificationParser.parse(content, request.format) MockGenerationContext(request, specification) } val generateNode by node<MockGenerationContext, MockGenerationContext>("generate") { ctx -> val prompt = promptBuilder.buildSpecWithDescriptionPrompt( ctx.specification, ctx.request.description, ctx.request.namespace, ctx.specification.format ) val response = llm.writeSession { appendPrompt { user(prompt) } requestLLM() } val mocks = aiModelService.parseModelResponse(textResponse, ctx.request.namespace, ...) ctx.copy(mocks = mocks) } val validateNode by node<MockGenerationContext, MockGenerationContext>("validate") { ctx -> val errors = ctx.mocks.flatMap { mockValidator.validate(it, ctx.specification).errors } ctx.copy(errors = errors) } val correctNode by node<MockGenerationContext, MockGenerationContext>("correct") { ctx -> val correctionPrompt = promptBuilder.buildCorrectionPrompt( invalidMocks = ..., namespace = ctx.request.namespace, specification = ctx.specification, format = ctx.specification.format ) val response = llm.writeSession { appendPrompt { user(correctionPrompt) } requestLLM() } // Parse corrected mocks, merge with valid ones from first pass } edge(nodeStart forwardTo setupNode) edge(setupNode forwardTo generateNode) edge(generateNode forwardTo validateNode onCondition { it.request.options.enableValidation }) edge(generateNode forwardTo nodeFinish onCondition { !it.request.options.enableValidation }) edge(validateNode forwardTo nodeFinish onCondition { it.errors.isEmpty() || it.attempt > maxRetries }) edge(validateNode forwardTo correctNode onCondition { it.errors.isNotEmpty() && it.attempt <= maxRetries }) edge(correctNode forwardTo validateNode)}Lesson 4: Watch Out for Bedrock’s Output Token Limit
I was generating mocks from a large GraphQL schema, and the output kept coming back as truncated JSON. I kept tweaking the JSON formatting instructions, thinking the model was producing malformed output. The correction loop retried, but the corrected output was also truncated.
The problem wasn’t the JSON part of the prompt — it was Bedrock’s
maxOutputTokens limit. The schema had so many queries that the mocks exceeded the token ceiling every time. The fix was asking the model to generate less by adjusting a completely different part of the prompt.Output token limits vary by model — check the Bedrock model documentation [8] and service quotas [9]. If your output is consistently truncated, retrying won’t help.
Lesson 5: Testing the Graph Without Spending Money
Koog’s
agents-test module lets you create a mock LLM executor, wire it into a real GraphAIAgent , and run the actual strategy graph with fake responses:val mockExecutor = getMockExecutor { mockLLMAnswer(validJsonResponse).asDefaultResponse}val agent = GraphAIAgent( inputType = strategy.inputType, outputType = strategy.outputType, promptExecutor = mockExecutor, agentConfig = agentConfig, strategy = strategy, toolRegistry = ToolRegistry.EMPTY)val result = agent.run(request)assertTrue(result.success)You can test happy paths, correction loops, parse failure recovery, and prompt content — all without AWS credentials or spending a cent. Tests run in milliseconds.
But there’s one thing these tests can’t tell you: are your prompts any good?
Lesson 6: Testing Your Prompts — Eval-Driven Development
You can’t unit test prompts. You need to run them against the real model and measure output quality. MockNest has a dedicated eval suite built with Dokimos [11] — a Kotlin LLM evaluation framework. Koog provides a
dokimos-koog integration module [12] that bridges Koog's prompt executor into Dokimos's judge interface via asJudge . The suite generates mocks from real specs, validates them, runs the correction loop, and uses Dokimos's LLMJudgeEvaluator for semantic correctness (an LLM-as-a-judge pattern). Since we use the Converse API, we get token usage back from every Bedrock call, which lets us track costs per scenario.The eval test loads scenarios from a dataset JSON, runs each one through the real Koog strategy graph against Bedrock, and then asks the LLM judge to score the output:
// Wire up Dokimos's LLM judge using Koog's executorval judgeLM = asJudge { prompt -> val judgeAgent = AIAgent( promptExecutor = judgeExecutor, llmModel = modelConfiguration.getModel(), systemPrompt = "You are an evaluation judge. Respond only with a numeric score.", toolRegistry = ToolRegistry.EMPTY ) judgeAgent.run(prompt)}
// Evaluate semantic correctness of generated mocksval evaluator = LLMJudgeEvaluator.builder() .name("Scenario Semantic Check") .criteria(scenario.semanticCheck) // e.g. "Output contains valid WireMock mappings for all GraphQL queries" .evaluationParams(listOf(EvalTestCaseParam.INPUT, EvalTestCaseParam.ACTUAL_OUTPUT)) .judge(judgeLM) .threshold(0.7) .build()val testCase = EvalTestCase.builder() .input(scenario.description) .actualOutput(mappingsJson) .expectedOutput("") .build()val result = evaluator.evaluate(testCase)// result.score(), result.success()Token usage is captured by decorating the
BedrockRuntimeClient — intercepting converse calls and recording input/output tokens from each response. This keeps cost tracking independent of Koog internals.The test is gated behind an environment variable so it never runs in CI by default:
BEDROCK_EVAL_ENABLED=true ./gradlew :software:infra:aws:generation:bedrockEvalExample summary table:
╔═══════════════════════════════════════════════════════════════════════════════════╗║ MULTI-PROTOCOL BEDROCK PROMPT EVAL SUMMARY ║╠═══════════════════════════════════════════════════════════════════════════════════╣║ Protocol │ Runs │ 1st-pass valid │ After retry │ Scenario pass │ Cost │ Latency ║╠═══════════════════════════════════════════════════════════════════════════════════╣║ GraphQL │ 15 │ 87% │ 98% │ 98% │ $0.004 │ 2.8s ║║ SOAP │ 15 │ 77% │ 97% │ 97% │ $0.009 │ 6.1s ║╚═══════════════════════════════════════════════════════════════════════════════════╝These tests have to be triggered manually because there’s a cost attached. This is a remocal testing [13][14] approach — local test code calling remote cloud services — which gives you real model feedback without deploying anything.
Wrapping Up
Koog and Bedrock work great together once you know what to watch for:
- Use IAM roles instead of access keys.
- Don’t set
credentialsProvideron the client — lazy resolution via the default chain avoids eager connections.
- If you need multiple LLM calls within the same session, use Koog’s strategy graph — it manages the session across nodes and keeps wiring separate from logic.
- Watch out for output token limits.
- Test the graph with
agents-test— fast, free, deterministic.
- Test your prompts with eval tests against real Bedrock using a remocal approach — trigger manually for anything with a cost attached.
- Keep Bedrock out of your business logic so the whole thing stays portable.
References
Series: Kotlin on AWS (1 article)
- 1Lessons Learned: Shipping a Koog + Bedrock AI Agent to Production on AWS Lambda This article
Enjoyed reading this content? Let the author know!
Your likes, comments, shares, and saves help creators reach more builders.
Loading recommendations
Loading article