Build and Test AWS Lambda Extensions Locally with LocalStack

Introduction

AWS Lambda extensions let you hook into the Lambda execution environment to run your own code alongside the function handler. They are useful for capturing telemetry, enforcing security policies, caching secrets, or integrating with observability tools, all without modifying the function code itself.

Building and testing extensions on AWS is slow. Each iteration involves packaging a layer, updating the function, invoking it, then waiting for the environment to shut down so you can inspect the results. That shutdown can take several minutes after you delete the function.

LocalStack makes this loop shorter. The execution environment shuts down within seconds of function deletion, and you can inspect S3, CloudWatch Logs, and other resources immediately. In this tutorial, we’ll build an audit-log extension that records every Lambda invocation and flushes the data to S3 when the environment shuts down.

How Lambda extensions work

Lambda extensions come in two forms: internal and external. Internal extensions run as part of the runtime process (wrapper scripts or language-specific hooks). External extensions run as independent processes in the same execution environment but separate from the runtime. They start during INIT, run alongside the function, and get notified when the environment is about to shut down. This tutorial focuses on external extensions.

The Extensions API lifecycle

An external extension goes through three phases:

1. INIT: The extension registers itself by calling POST /2020-01-01/extension/register with the list of events it wants to receive (INVOKE, SHUTDOWN, or both). The registration response includes a Lambda-Extension-Identifier header used for all subsequent API calls.
2. INVOKE: The extension enters a polling loop, calling GET /2020-01-01/extension/event/next. This blocking call returns when the next event is available. The payload includes the requestId, function ARN, deadline timestamp, and X-Ray tracing information.
3. SHUTDOWN: When the execution environment is about to be destroyed, the extension receives a SHUTDOWN event with a shutdownReason and deadlineMs. This is the extension’s last chance to flush buffers, close connections, or write final output.

The extension process runs independently of the function handler. They share the same filesystem and network but do not block each other.

Extension packaging

Extensions are packaged as Lambda layers. The layer zip must contain an extensions/ directory at the root, and each file inside that directory becomes an extension process. The file must be executable. Lambda unpacks the layer into /opt/extensions/ and starts each one as a separate process during INIT.

For a Python extension, the zip structure looks like:

extensions/
└── audit-log      # executable Python script with a shebang

The filename becomes the extension name used in logs and error reporting.

This tutorial uses the Extensions API lifecycle events (INVOKE and SHUTDOWN). Some third-party extensions depend on the Lambda Telemetry API instead, so check LocalStack support before testing telemetry-focused extensions.

Prerequisites

Before starting, make sure you have the following installed:

• Docker
• LocalStack CLI with a valid LocalStack Auth Token and access to Lambda layers. This tutorial publishes the extension as a Lambda layer, which requires a LocalStack plan that includes Lambda layer support.
• AWS CLI with the awslocal wrapper script
• Python 3.12+ (for the Lambda function and extension code)
• zip

Start LocalStack with your Auth Token configured:

export LOCALSTACK_AUTH_TOKEN=your-auth-token
localstack start

Wait for LocalStack to be ready:

localstack wait -t 60

Step 1: Set up the project structure

Create a new directory for the project:

mkdir -p lambda-extensions-sample/extension/extensions
mkdir -p lambda-extensions-sample/function
cd lambda-extensions-sample

The tutorial is self-contained: the commands below create the complete sample from an empty directory, so you can keep the resulting lambda-extensions-sample directory as a reusable local example.

The project will have this structure when we are done:

lambda-extensions-sample/
├── extension/
│   └── extensions/
│       └── audit-log          # the extension script (executable)
└── function/
    └── handler.py             # the Lambda function handler

The extension/extensions/ nesting is intentional. When we zip the extension/ directory, the archive will have extensions/audit-log at the root, which is what Lambda expects when it unpacks the layer into /opt/.

Step 2: Write the extension

Create the file extension/extensions/audit-log with the following content:

#!/usr/bin/env python3
"""
audit-log: A Lambda External Extension that tracks every invocation and writes
an audit report to S3 when the execution environment shuts down.
"""

import json
import os
import sys
import time
import urllib.request

EXTENSION_NAME = os.path.basename(__file__)
RUNTIME_API = os.environ["AWS_LAMBDA_RUNTIME_API"]
BASE_URL = f"http://{RUNTIME_API}/2020-01-01/extension"

AUDIT_BUCKET = os.environ.get("AUDIT_BUCKET", "")
FUNCTION_NAME = os.environ.get("AWS_LAMBDA_FUNCTION_NAME", "unknown")

AWS_LAMBDA_RUNTIME_API provides the local HTTP endpoint for the Extensions API, AWS_LAMBDA_FUNCTION_NAME identifies which function this environment belongs to, and AUDIT_BUCKET is a custom environment variable we’ll set when creating the function.

Next, add the Extensions API helper functions:

def register(events: list[str]) -> str:
    """Register with the Extensions API and return the extension identifier."""
    url = f"{BASE_URL}/register"
    data = json.dumps({"events": events}).encode()
    req = urllib.request.Request(
        url,
        data=data,
        headers={"Lambda-Extension-Name": EXTENSION_NAME},
        method="POST",
    )
    with urllib.request.urlopen(req) as resp:
        extension_id = resp.headers["Lambda-Extension-Identifier"]
        body = json.loads(resp.read())
        log(f"Registered – id={extension_id}, function={body.get('functionName')}")
        return extension_id


def next_event(extension_id: str) -> dict:
    """Block until the next INVOKE or SHUTDOWN event arrives."""
    url = f"{BASE_URL}/event/next"
    req = urllib.request.Request(
        url,
        headers={"Lambda-Extension-Identifier": extension_id},
        method="GET",
    )
    with urllib.request.urlopen(req) as resp:
        return json.loads(resp.read())

These two functions are the core of the Extensions API interaction. register() tells Lambda which events we care about and returns an identifier for subsequent requests. next_event() is a long-polling call that blocks until Lambda sends an invocation or shutdown event.

Add a function to report errors if something goes wrong:

def report_error(extension_id: str, error_type: str, message: str):
    """Report an error to the Extensions API before exiting."""
    url = f"{BASE_URL}/exit/error"
    data = json.dumps(
        {"errorType": error_type, "errorMessage": message, "stackTrace": []}
    ).encode()
    req = urllib.request.Request(
        url,
        data=data,
        headers={
            "Lambda-Extension-Identifier": extension_id,
            "Lambda-Extension-Function-Error-Type": error_type,
        },
        method="POST",
    )
    with urllib.request.urlopen(req) as resp:
        resp.read()

Now add the S3 upload logic. Since boto3 ships with every Python Lambda runtime under /var/runtime, the extension can import it directly:

sys.path.insert(0, "/var/runtime")
import boto3


def get_s3_client():
    """Return an S3 client, respecting AWS_ENDPOINT_URL for LocalStack."""
    kwargs = {}
    endpoint = os.environ.get("AWS_ENDPOINT_URL")
    if endpoint:
        kwargs["endpoint_url"] = endpoint
    return boto3.client("s3", **kwargs)


def flush_audit_log(records: list[dict], shutdown_reason: str):
    """Write the collected audit records to S3 as a single JSON file."""
    if not AUDIT_BUCKET:
        log("AUDIT_BUCKET not set – skipping S3 upload")
        return

    report = {
        "function_name": FUNCTION_NAME,
        "extension": EXTENSION_NAME,
        "total_invocations": len(records),
        "shutdown_reason": shutdown_reason,
        "generated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "invocations": records,
    }

    key = f"audit-logs/{FUNCTION_NAME}/{int(time.time())}.json"
    s3 = get_s3_client()
    s3.put_object(
        Bucket=AUDIT_BUCKET,
        Key=key,
        Body=json.dumps(report, indent=2),
        ContentType="application/json",
    )
    log(f"Audit log written to s3://{AUDIT_BUCKET}/{key} ({len(records)} invocations)")

get_s3_client() checks for AWS_ENDPOINT_URL, which LocalStack sets automatically inside the Lambda execution environment. On real AWS this variable is absent, so the client uses the standard S3 endpoint. No code changes needed between LocalStack and AWS.

Finally, add the main loop and a logging helper:

def log(msg: str):
    print(f"[{EXTENSION_NAME}] {msg}", flush=True)


def main():
    init_start = time.time()
    extension_id = register(["INVOKE", "SHUTDOWN"])
    init_duration_ms = round((time.time() - init_start) * 1000)
    log(f"Init completed in {init_duration_ms}ms")

    audit_records = []

    while True:
        event = next_event(extension_id)
        event_type = event.get("eventType")

        if event_type == "INVOKE":
            record = {
                "request_id": event.get("requestId"),
                "invoked_function_arn": event.get("invokedFunctionArn"),
                "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
                "deadline_ms": event.get("deadlineMs"),
                "trace_id": event.get("tracing", {}).get("value", ""),
                "invocation_number": len(audit_records) + 1,
            }
            audit_records.append(record)
            log(f"INVOKE #{record['invocation_number']} – request_id={record['request_id']}")

        elif event_type == "SHUTDOWN":
            shutdown_reason = event.get("shutdownReason", "unknown")
            log(f"SHUTDOWN received – reason={shutdown_reason}")
            flush_audit_log(audit_records, shutdown_reason)
            log("Exiting cleanly")
            return


if __name__ == "__main__":
    try:
        main()
    except Exception as exc:
        log(f"Fatal error: {exc}")
        sys.exit(1)

The main loop is straightforward: register, then poll for events. Each INVOKE event gets recorded with its metadata. When SHUTDOWN arrives, the extension flushes everything to S3 and exits.

Make the extension executable:

chmod +x extension/extensions/audit-log

Step 3: Write the Lambda function

Create function/handler.py. This is a simple order-processing function that gives the extension something to observe:

import json
import time
import uuid


def handler(event, context):
    action = event.get("action", "unknown")

    if action == "place_order":
        order = {
            "order_id": str(uuid.uuid4()),
            "items": event.get("items", []),
            "total": sum(item.get("price", 0) for item in event.get("items", [])),
            "status": "confirmed",
            "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        }
        return {"statusCode": 200, "body": json.dumps(order)}

    elif action == "get_status":
        return {
            "statusCode": 200,
            "body": json.dumps(
                {"order_id": event.get("order_id", ""), "status": "processing"}
            ),
        }

    return {
        "statusCode": 400,
        "body": json.dumps({"error": f"Unknown action: {action}"}),
    }

The audit-log extension works with any Lambda function since it observes lifecycle events, not function payloads. We are using this handler so we have something realistic to invoke.

Step 4: Package and deploy to LocalStack

Create an S3 bucket for the audit logs:

awslocal s3 mb s3://audit-logs-bucket

Package the extension. The zip must have the extensions/ directory at the root:

cd extension
zip -r /tmp/audit-log-extension.zip extensions/
cd ..

Publish it as a Lambda layer:

awslocal lambda publish-layer-version \
    --layer-name audit-log-extension \
    --zip-file fileb:///tmp/audit-log-extension.zip \
    --compatible-runtimes python3.12 python3.13

You should see output that includes the layer ARN:

{
    "LayerArn": "arn:aws:lambda:us-east-1:000000000000:layer:audit-log-extension",
    "LayerVersionArn": "arn:aws:lambda:us-east-1:000000000000:layer:audit-log-extension:1",
    "Version": 1,
    "CompatibleRuntimes": ["python3.12", "python3.13"]
}

Package the function code:

cd function
zip /tmp/order-processor.zip handler.py
cd ..

Create an IAM execution role:

awslocal iam create-role \
    --role-name lambda-extension-role \
    --assume-role-policy-document '{
        "Version": "2012-10-17",
        "Statement": [{
            "Effect": "Allow",
            "Principal": {"Service": "lambda.amazonaws.com"},
            "Action": "sts:AssumeRole"
        }]
    }'

Create the Lambda function with the extension layer attached:

awslocal lambda create-function \
    --function-name order-processor \
    --runtime python3.12 \
    --handler handler.handler \
    --role arn:aws:iam::000000000000:role/lambda-extension-role \
    --zip-file fileb:///tmp/order-processor.zip \
    --timeout 30 \
    --layers arn:aws:lambda:us-east-1:000000000000:layer:audit-log-extension:1 \
    --environment 'Variables={AUDIT_BUCKET=audit-logs-bucket}'

The --layers flag is what connects the extension to the function. When Lambda creates the execution environment, it unpacks this layer into /opt/, finds the executable at /opt/extensions/audit-log, and starts it as a separate process.

Wait for the function to become active:

awslocal lambda wait function-active-v2 --function-name order-processor

Step 5: Invoke the function and observe the extension

With the function deployed, let’s invoke it a few times. Each invocation will trigger an INVOKE event that the extension captures.

Place an order:

awslocal lambda invoke \
    --function-name order-processor \
    --cli-binary-format raw-in-base64-out \
    --payload '{"action":"place_order","items":[{"name":"Laptop","price":999},{"name":"Mouse","price":29}]}' \
    /tmp/response.json

Check the response:

cat /tmp/response.json | python3 -m json.tool

The output should look like this:

{
    "statusCode": 200,
    "body": "{\"order_id\": \"4cef8fd4-94ac-43b7-a979-30f85bfb8c9b\", \"items\": [{\"name\": \"Laptop\", \"price\": 999}, {\"name\": \"Mouse\", \"price\": 29}], \"total\": 1028, \"status\": \"confirmed\", \"timestamp\": \"2026-04-03T11:02:05Z\"}"
}

Run two more invocations:

awslocal lambda invoke \
    --function-name order-processor \
    --cli-binary-format raw-in-base64-out \
    --payload '{"action":"get_status","order_id":"abc-123"}' \
    /tmp/response.json

awslocal lambda invoke \
    --function-name order-processor \
    --cli-binary-format raw-in-base64-out \
    --payload '{"action":"place_order","items":[{"name":"Keyboard","price":75}]}' \
    /tmp/response.json

The extension has recorded all three invocations in memory. Verify this by checking CloudWatch Logs:

awslocal logs filter-log-events \
    --log-group-name /aws/lambda/order-processor \
    --query 'events[].message' \
    --output text

You should see the extension’s log lines interleaved with the standard Lambda entries:

[audit-log] Registered – id=9b5106fe-e31e-402b-bf44-d8a4a1af3e75, function=order-processor
[audit-log] Init completed in 103ms
START RequestId: fb5fc9b2-... Version: $LATEST
[audit-log] INVOKE #1 – request_id=fb5fc9b2-...
END RequestId: fb5fc9b2-...
REPORT RequestId: fb5fc9b2-...  Duration: 17.24 ms ...
START RequestId: 42048220-... Version: $LATEST
[audit-log] INVOKE #2 – request_id=42048220-...
END RequestId: 42048220-...
REPORT RequestId: 42048220-...  Duration: 16.46 ms ...
START RequestId: 01f24d06-... Version: $LATEST
[audit-log] INVOKE #3 – request_id=01f24d06-...
END RequestId: 01f24d06-...
REPORT RequestId: 01f24d06-...  Duration: 28.59 ms ...

The [audit-log] lines come from the extension process. The START/END/REPORT lines come from the Lambda runtime, running as independent processes in the same execution environment.

Step 6: Trigger shutdown and read the audit log

The extension flushes its data to S3 only on SHUTDOWN. Delete the function to trigger it:

awslocal lambda delete-function --function-name order-processor

On LocalStack, the execution environment shuts down within seconds. On AWS, this can take several minutes. List the audit log objects:

awslocal s3 ls s3://audit-logs-bucket/audit-logs/order-processor/

Download and inspect the audit log:

awslocal s3 cp s3://audit-logs-bucket/audit-logs/order-processor/ /tmp/audit/ --recursive
cat /tmp/audit/*.json

The output is the full audit report written by the extension during shutdown:

{
    "function_name": "order-processor",
    "extension": "audit-log",
    "total_invocations": 3,
    "shutdown_reason": "SandboxTerminated",
    "generated_at": "2026-04-03T11:02:23Z",
    "invocations": [
        {
            "request_id": "fb5fc9b2-9f7a-4853-a16a-392fb3c814fb",
            "invoked_function_arn": "arn:aws:lambda:us-east-1:000000000000:function:order-processor",
            "timestamp": "2026-04-03T11:02:05Z",
            "deadline_ms": 1775214155008,
            "trace_id": "Root=1-69cf9e29-c294024b52ce6b2a5be5a6b7;Parent=47d8b4750256da8d;Sampled=1",
            "invocation_number": 1
        },
        {
            "request_id": "42048220-e1b4-424a-9ff9-1ddb7eaf44fb",
            "invoked_function_arn": "arn:aws:lambda:us-east-1:000000000000:function:order-processor",
            "timestamp": "2026-04-03T11:02:10Z",
            "deadline_ms": 1775214160841,
            "trace_id": "Root=1-69cf9e32-6dfe53979ef565bcfc239958;Parent=3748759ab7bdb4d1;Sampled=1",
            "invocation_number": 2
        },
        {
            "request_id": "01f24d06-ec1a-4784-b984-ab14a75a8d26",
            "invoked_function_arn": "arn:aws:lambda:us-east-1:000000000000:function:order-processor",
            "timestamp": "2026-04-03T11:02:16Z",
            "deadline_ms": 1775214166289,
            "trace_id": "Root=1-69cf9e38-44fddba7a7907fa339ce31fb;Parent=4db19c1d4d3ac04c;Sampled=1",
            "invocation_number": 3
        }
    ]
}

Each entry captures the request ID (matching the Lambda START log), function ARN, X-Ray trace ID, and deadline timestamp. The shutdown_reason is SandboxTerminated since the function was deleted. Other possible reasons are timeout and failure.

Summary

In this tutorial, we built a Lambda external extension that:

• Registers with the Extensions API for INVOKE and SHUTDOWN events
• Records metadata for every function invocation
• Flushes a structured audit report to S3 on shutdown
• Uses boto3 from the Lambda runtime with transparent endpoint injection for LocalStack

LocalStack supports the Lambda Extensions API flow used in this tutorial. The execution environment shuts down in seconds instead of minutes, and you can iterate on extension code without redeploying to a cloud account. For extensions that rely on other Lambda extension surfaces, such as the Telemetry API, confirm current LocalStack support before using them in your local workflow.