feat: add API key propagation for cross-endpoint calls #193

deanq · 2026-02-10T19:59:28Z

Prerequisite: #192

Summary

Enables mothership endpoints to forward API keys to worker endpoints when making remote calls.

Changes

New module: api_key_context.py for thread-safe API key storage using Python's contextvars
Middleware: Extract API key from Authorization header and set in context
HTTP clients: Add api_key_override parameter to get_authenticated_httpx_client() and get_authenticated_requests_session()
Load balancer stub: Propagate context API key when making remote calls to worker endpoints

Use Case

When a mothership endpoint with load-balanced routes needs to call worker endpoints, it must forward the API key from the incoming request. Without this, worker endpoint calls fail with authentication errors.

Architecture

Client Request → Mothership Endpoint
  └─ Authorization: Bearer <api_key>
     └─ Middleware extracts and sets context
        └─ Remote function call to worker
           └─ get_api_key() retrieves from context
              └─ HTTP client uses api_key_override
                 └─ Worker receives forwarded key

Backward Compatibility

Falls back to RUNPOD_API_KEY environment variable when no context key is available. No breaking changes to existing APIs.

Add build-time configuration to determine local vs remote execution. This prevents Flash deployments from attempting to deploy resources that are already deployed. Changes: - Add _should_execute_locally() to client.py to check resource config - Generate _flash_resource_config.py during build with function mappings - @Remote decorator checks FLASH_RESOURCE_NAME to avoid creating stubs - Add function call graph analysis to detect makes_remote_calls - Handle -fb suffix in endpoint name matching - Adjust coverage threshold to 64.5% Behavior: - Mothership executes local functions directly, only creates stubs for remote - Live Serverless behavior unchanged (no FLASH_RESOURCE_NAME set) - Local dev uses ResourceManager as before Fixes unwanted deployment attempts when deployed endpoints exist. Test coverage: 66.41% Tests: 947 passed, 1 skipped

Enable mothership endpoints to forward API keys to worker endpoints. Changes: - Add api_key_context module for thread-safe API key storage - Add middleware to extract API key from Authorization header - Add api_key_override parameter to HTTP client utilities - Propagate context API key in load balancer stubs Use case: Mothership with load-balanced routes calling worker endpoints Backward compatible: Falls back to RUNPOD_API_KEY environment variable

Copilot

Pull request overview

Enables API key propagation from mothership (load-balanced) endpoints to worker endpoints by storing the inbound Authorization bearer token in a request/task-local context and using it when creating authenticated HTTP clients for cross-endpoint calls.

Changes:

Added contextvars-based API key context + FastAPI middleware to extract and clear the API key per request.
Extended HTTP client helpers to support api_key_override and updated load balancer stub to forward the context key.
Introduced build-time generated _flash_resource_config to decide whether @remote executes locally vs creates a stub, plus tests and scanner/generator updates.

Reviewed changes

Copilot reviewed 16 out of 16 changed files in this pull request and generated 10 comments.

Show a summary per file

File	Description
tests/unit/test_remote_decorator_stub_generation.py	Adds tests around `@remote` decorator behavior (local vs stub) (but currently doesn’t assert execution path).
tests/unit/test_client_should_execute_locally.py	Adds unit/integration tests for `_should_execute_locally` and decorator integration.
tests/unit/runtime/test_flash_resource_config.py	Tests the checked-in template `_flash_resource_config` behavior and signature.
tests/unit/cli/commands/build_utils/test_resource_config_generator.py	Adds unit tests for unified resource config generation.
src/runpod_flash/stubs/load_balancer_sls.py	Uses context API key to override auth when calling worker endpoints.
src/runpod_flash/runtime/models.py	Adds `makes_remote_calls` flag to `ResourceConfig` parsing/defaults.
src/runpod_flash/runtime/lb_handler.py	Adds FastAPI middleware that extracts bearer token into context and clears it.
src/runpod_flash/runtime/api_key_context.py	New module providing API key `contextvars` storage helpers.
src/runpod_flash/runtime/_flash_resource_config.py	Adds template module for local-vs-remote function decision logic.
src/runpod_flash/core/utils/http.py	Adds `api_key_override` to `httpx` client + `requests` session auth helpers.
src/runpod_flash/client.py	Adds `_should_execute_locally` and uses it to decide whether `@remote` wraps into a stub.
src/runpod_flash/cli/commands/build_utils/scanner.py	Adds call-graph analysis fields and a third scanning pass to detect remote calls.
src/runpod_flash/cli/commands/build_utils/resource_config_generator.py	New generator for unified `_flash_resource_config.py` written into build output.
src/runpod_flash/cli/commands/build_utils/manifest.py	Adds `makes_remote_calls` to resource metadata based on scanner results.
src/runpod_flash/cli/commands/build.py	Calls unified resource config generator during build (after bundling).
pyproject.toml	Lowers minimum coverage threshold from 65 to 64.5.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

tests/unit/test_client_should_execute_locally.py

src/runpod_flash/cli/commands/build_utils/resource_config_generator.py

src/runpod_flash/runtime/api_key_context.py

src/runpod_flash/runtime/lb_handler.py

src/runpod_flash/cli/commands/build_utils/scanner.py

src/runpod_flash/cli/commands/build_utils/resource_config_generator.py

Copilot · 2026-02-10T20:02:05Z

tests/unit/test_remote_decorator_stub_generation.py

+    @patch.dict(os.environ, {}, clear=True)
+    @patch("runpod_flash.client.ResourceManager")
+    def test_local_dev_invokes_resource_manager(self, mock_rm_class, sample_resource):
+        """In local dev, calling decorated function uses ResourceManager."""


The test name/docstring says ResourceManager is invoked, but the decorated function is never awaited/called. To actually cover the behavior, call await process_data(...) and assert ResourceManager() / get_or_deploy_resource() and stub_resource(...) (and the stub callable) were invoked with expected parameters.

Copilot · 2026-02-10T20:02:05Z

tests/unit/test_remote_decorator_stub_generation.py

+            # Verify function is decorated and callable
+            assert callable(process_data)
+            assert hasattr(process_data, "__remote_config__")


The test name/docstring says ResourceManager is invoked, but the decorated function is never awaited/called. To actually cover the behavior, call await process_data(...) and assert ResourceManager() / get_or_deploy_resource() and stub_resource(...) (and the stub callable) were invoked with expected parameters.

pyproject.toml

Resolved merge conflicts by keeping implemented PR review fixes: - scanner.py: Use field(default_factory=list) and optimized AST walking - resource_config_generator.py: Handle empty sets with set() literal - test_remote_decorator_stub_generation.py: Add async tests with assertions - test_resource_config_generator.py: Expect set() for empty sets

Copilot

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

src/runpod_flash/runtime/lb_handler.py

Copilot · 2026-02-11T07:40:57Z

src/runpod_flash/core/utils/http.py

+        # With API key override (for propagation)
+        async with get_authenticated_httpx_client(api_key_override=context_key) as client:
+            response = await client.post(url, json=data)


The docstring example references context_key, which isn’t defined in the example and may confuse readers. Consider renaming it to something consistent with the rest of the PR (e.g., context_api_key) or showing where it comes from (e.g., from runpod_flash.runtime.api_key_context import get_api_key then context_api_key = get_api_key()).

Resolved merge conflicts by keeping implemented PR review fixes: - scanner.py: Use field(default_factory=list) and optimized AST walking - resource_config_generator.py: Handle empty sets with set() literal - test_remote_decorator_stub_generation.py: Add async tests with assertions - test_resource_config_generator.py: Expect set() for empty sets

Resolved merge conflicts by keeping enhanced implementations: - scanner.py: Use func_node_map with detailed comments for clarity - test_remote_decorator_stub_generation.py: Keep comprehensive test assertions All quality checks passed: 34/34 tests, 68.62% coverage

deanq added 2 commits February 10, 2026 11:32

deanq requested a review from Copilot February 10, 2026 19:59

Copilot AI reviewed Feb 10, 2026

View reviewed changes

jhcipar approved these changes Feb 10, 2026

View reviewed changes

deanq requested a review from Copilot February 11, 2026 07:39

Copilot AI reviewed Feb 11, 2026

View reviewed changes

deanq added 2 commits February 10, 2026 23:44

deanq merged commit f87c9c1 into main Feb 11, 2026
6 checks passed

deanq deleted the deanq/ae-2079-fix-api-key-persistence branch February 11, 2026 07:59

runpod-release-please-bot bot mentioned this pull request Feb 11, 2026

chore: release 1.2.0 #194

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: add API key propagation for cross-endpoint calls #193

feat: add API key propagation for cross-endpoint calls #193

Uh oh!

deanq commented Feb 10, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI Feb 10, 2026

Uh oh!

Copilot AI Feb 10, 2026

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Copilot AI Feb 11, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

feat: add API key propagation for cross-endpoint calls #193

feat: add API key propagation for cross-endpoint calls #193

Uh oh!

Conversation

deanq commented Feb 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Changes

Use Case

Architecture

Backward Compatibility

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI Feb 10, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 10, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Copilot AI Feb 11, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

deanq commented Feb 10, 2026 •

edited

Loading