-
Notifications
You must be signed in to change notification settings - Fork 125
feat: add aibridge-proxy module for AI Bridge Proxy workspace setup #721
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Open
ssncferreira
wants to merge
5
commits into
main
Choose a base branch
from
ssncf/feat-aibridge-proxy-module
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
+713
−0
Open
Changes from all commits
Commits
Show all changes
5 commits
Select commit
Hold shift + click to select a range
05c41f7
feat: add aibridge-proxy module for AI Bridge Proxy workspace setup
ssncferreira bd29232
feat: signal startup coordination on setup completion
ssncferreira 1939192
docs: update README
ssncferreira 50ecb11
chore: fix startup complete coordination
ssncferreira 9ab3b40
chore: improve documentation based on comments
ssncferreira File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,89 @@ | ||
| --- | ||
| display_name: AI Bridge Proxy | ||
| description: Configure a workspace to route AI tool traffic through AI Bridge via AI Bridge Proxy. | ||
| icon: ../../../../.icons/coder.svg | ||
| verified: true | ||
| tags: [helper, aibridge] | ||
| --- | ||
|
|
||
| # AI Bridge Proxy | ||
|
|
||
| This module configures a Coder workspace to use [AI Bridge Proxy](https://coder.com/docs/ai-coder/ai-bridge/ai-bridge-proxy). | ||
| It downloads the proxy's CA certificate from the Coder deployment and provides Terraform outputs (`proxy_auth_url` and `cert_path`) that tool-specific modules can use to route their traffic through the proxy. | ||
|
|
||
| ```tf | ||
| module "aibridge-proxy" { | ||
| source = "registry.coder.com/coder/aibridge-proxy/coder" | ||
| version = "1.0.0" | ||
| agent_id = coder_agent.main.id | ||
| proxy_url = "https://aiproxy.example.com" | ||
| } | ||
| ``` | ||
|
|
||
| > [!NOTE] | ||
| > AI Bridge Proxy is a Premium Coder feature that requires [AI Governance Add-On](https://coder.com/docs/ai-coder/ai-governance). | ||
| > See the [AI Bridge Proxy setup guide](https://coder.com/docs/ai-coder/ai-bridge/ai-bridge-proxy/setup) for details on configuring the proxy on your Coder deployment. | ||
|
|
||
| ## How it works | ||
|
|
||
| [AI Bridge Proxy](https://coder.com/docs/ai-coder/ai-bridge/ai-bridge-proxy) is an HTTP proxy that intercepts traffic to AI providers and forwards it through [AI Bridge](https://coder.com/docs/ai-coder/ai-bridge), enabling centralized LLM management, governance, and cost tracking. | ||
| Any process with the proxy environment variables set will route **all** its traffic through the proxy. | ||
|
|
||
| This module **does not** set proxy environment variables globally on the workspace. | ||
| Instead, it provides Terraform outputs (`proxy_auth_url` and `cert_path`) that tool-specific modules consume to configure proxy routing. | ||
| See the [Copilot module](https://registry.coder.com/modules/coder-labs/copilot) for a working integration example. | ||
|
|
||
| It is recommended that tool modules scope the proxy environment variables to their own process rather than setting them globally on the workspace, to avoid routing unnecessary traffic through the proxy. | ||
|
|
||
| > [!WARNING] | ||
| > If the setup script fails (e.g. the proxy is unreachable), the workspace will still start but the agent will report a startup script error. | ||
| > Tools that depend on the proxy will not work until the issue is resolved. Check the workspace build logs for details. | ||
|
|
||
| ## Startup Coordination | ||
|
|
||
| When used with tool-specific modules (e.g. [Copilot](https://registry.coder.com/modules/coder-labs/copilot)), | ||
| the setup script signals completion via [`coder exp sync`](https://coder.com/docs/admin/templates/startup-coordination) so dependent modules can wait for the `aibridge-proxy` module to complete before starting. | ||
|
|
||
| Dependent modules are unblocked once the setup script finishes, regardless of success or failure. | ||
| If the setup fails, dependent modules are expected to detect the failure and handle the error accordingly. | ||
|
|
||
| To enable startup coordination, set `CODER_AGENT_SOCKET_SERVER_ENABLED=true` in the workspace container environment: | ||
|
|
||
| ```hcl | ||
| env = [ | ||
| "CODER_AGENT_TOKEN=${coder_agent.main.token}", | ||
| "CODER_AGENT_SOCKET_SERVER_ENABLED=true", | ||
| ] | ||
| ``` | ||
|
|
||
| > [!NOTE] | ||
| > [Startup coordination](https://coder.com/docs/admin/templates/startup-coordination) requires Coder >= v2.30. | ||
| > Without it, the sync calls are skipped gracefully but dependent modules may fail to start if the `aibridge-proxy` setup has not completed in time. | ||
|
|
||
| ## Examples | ||
|
|
||
| ### Custom certificate path | ||
|
|
||
| ```tf | ||
| module "aibridge-proxy" { | ||
| source = "registry.coder.com/coder/aibridge-proxy/coder" | ||
| version = "1.0.0" | ||
| agent_id = coder_agent.main.id | ||
| proxy_url = "https://aiproxy.example.com" | ||
| cert_path = "/home/coder/.certs/aibridge-proxy-ca.pem" | ||
| } | ||
| ``` | ||
|
|
||
| ### Proxy with custom port | ||
|
|
||
| For deployments where the proxy is accessed directly on a configured port. | ||
| See [security considerations](https://coder.com/docs/ai-coder/ai-bridge/ai-bridge-proxy/setup#security-considerations) for network access guidelines. | ||
|
|
||
| ```tf | ||
| module "aibridge-proxy" { | ||
| source = "registry.coder.com/coder/aibridge-proxy/coder" | ||
| version = "1.0.0" | ||
| agent_id = coder_agent.main.id | ||
| proxy_url = "http://internal-proxy:8888" | ||
| } | ||
| ``` | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,254 @@ | ||
| import { serve } from "bun"; | ||
| import { | ||
| afterEach, | ||
| beforeAll, | ||
| describe, | ||
| expect, | ||
| it, | ||
| setDefaultTimeout, | ||
| } from "bun:test"; | ||
| import { | ||
| execContainer, | ||
| findResourceInstance, | ||
| removeContainer, | ||
| runContainer, | ||
| runTerraformApply, | ||
| runTerraformInit, | ||
| testRequiredVariables, | ||
| } from "~test"; | ||
|
|
||
| let cleanupFunctions: (() => Promise<void>)[] = []; | ||
| const registerCleanup = (cleanup: () => Promise<void>) => { | ||
| cleanupFunctions.push(cleanup); | ||
| }; | ||
| afterEach(async () => { | ||
| const cleanupFnsCopy = cleanupFunctions.slice().reverse(); | ||
| cleanupFunctions = []; | ||
| for (const cleanup of cleanupFnsCopy) { | ||
| try { | ||
| await cleanup(); | ||
| } catch (error) { | ||
| console.error("Error during cleanup:", error); | ||
| } | ||
| } | ||
| }); | ||
|
|
||
| const FAKE_CERT = | ||
| "-----BEGIN CERTIFICATE-----\nMIIBfakecert\n-----END CERTIFICATE-----\n"; | ||
|
|
||
| // Runs terraform apply to render the setup script, then starts a Docker | ||
| // container where we can execute it against a mock server. | ||
| const setupContainer = async (vars: Record<string, string> = {}) => { | ||
| const state = await runTerraformApply(import.meta.dir, { | ||
| agent_id: "foo", | ||
| proxy_url: "https://aiproxy.example.com", | ||
| ...vars, | ||
| }); | ||
| const instance = findResourceInstance(state, "coder_script"); | ||
| const id = await runContainer("lorello/alpine-bash"); | ||
|
|
||
| registerCleanup(async () => { | ||
| await removeContainer(id); | ||
| }); | ||
|
|
||
| return { id, instance }; | ||
| }; | ||
|
|
||
| // Starts a mock HTTP server that simulates the Coder API certificate endpoint. | ||
| // Returns the server and its base URL. | ||
| const setupServer = (handler: (req: Request) => Response) => { | ||
| const server = serve({ | ||
| fetch: handler, | ||
| port: 0, | ||
| }); | ||
| registerCleanup(async () => { | ||
| server.stop(); | ||
| }); | ||
| return { | ||
| server, | ||
| // Base URL without trailing slash | ||
| url: server.url.toString().slice(0, -1), | ||
| }; | ||
| }; | ||
|
|
||
| setDefaultTimeout(30 * 1000); | ||
|
|
||
| describe("aibridge-proxy", () => { | ||
| beforeAll(async () => { | ||
| await runTerraformInit(import.meta.dir); | ||
| }); | ||
|
|
||
| // Verify that agent_id and proxy_url are required. | ||
| testRequiredVariables(import.meta.dir, { | ||
| agent_id: "foo", | ||
| proxy_url: "https://aiproxy.example.com", | ||
| }); | ||
|
|
||
| it("downloads the CA certificate successfully", async () => { | ||
| let receivedToken = ""; | ||
| const { url } = setupServer((req) => { | ||
| const reqUrl = new URL(req.url); | ||
| if (reqUrl.pathname === "/api/v2/aibridge/proxy/ca-cert.pem") { | ||
| receivedToken = req.headers.get("Coder-Session-Token") || ""; | ||
| return new Response(FAKE_CERT, { | ||
| status: 200, | ||
| headers: { "Content-Type": "application/x-pem-file" }, | ||
| }); | ||
| } | ||
| return new Response("not found", { status: 404 }); | ||
| }); | ||
|
|
||
| const { id, instance } = await setupContainer(); | ||
|
|
||
| // Override ACCESS_URL and SESSION_TOKEN at runtime to point at the mock server. | ||
| const exec = await execContainer(id, [ | ||
| "env", | ||
| `ACCESS_URL=${url}`, | ||
| "SESSION_TOKEN=test-session-token-123", | ||
| "bash", | ||
| "-c", | ||
| instance.script, | ||
| ]); | ||
| expect(exec.exitCode).toBe(0); | ||
| expect(exec.stdout).toContain( | ||
| "AI Bridge Proxy CA certificate saved to /tmp/aibridge-proxy/ca-cert.pem", | ||
| ); | ||
|
|
||
| // Verify the cert was written to the default path. | ||
| const certContent = await execContainer(id, [ | ||
| "cat", | ||
| "/tmp/aibridge-proxy/ca-cert.pem", | ||
| ]); | ||
| expect(certContent.stdout).toContain("BEGIN CERTIFICATE"); | ||
|
|
||
| // Verify the session token was sent in the request header. | ||
| expect(receivedToken).toBe("test-session-token-123"); | ||
| }); | ||
|
|
||
| it("fails when the server is unreachable", async () => { | ||
| const { id, instance } = await setupContainer(); | ||
|
|
||
| // Port 9999 has nothing listening, so curl will fail to connect. | ||
| const exec = await execContainer(id, [ | ||
| "env", | ||
| "ACCESS_URL=http://localhost:9999", | ||
| "SESSION_TOKEN=mock-token", | ||
| "bash", | ||
| "-c", | ||
| instance.script, | ||
| ]); | ||
| expect(exec.exitCode).not.toBe(0); | ||
| expect(exec.stdout).toContain( | ||
| "AI Bridge Proxy setup failed: could not connect to", | ||
| ); | ||
| }); | ||
|
|
||
| it("fails when the server returns a non-200 status", async () => { | ||
| const { url } = setupServer(() => { | ||
| return new Response("not found", { status: 404 }); | ||
| }); | ||
|
|
||
| const { id, instance } = await setupContainer(); | ||
|
|
||
| const exec = await execContainer(id, [ | ||
| "env", | ||
| `ACCESS_URL=${url}`, | ||
| "SESSION_TOKEN=mock-token", | ||
| "bash", | ||
| "-c", | ||
| instance.script, | ||
| ]); | ||
| expect(exec.exitCode).not.toBe(0); | ||
| expect(exec.stdout).toContain( | ||
| "AI Bridge Proxy setup failed: unexpected response", | ||
| ); | ||
| }); | ||
|
|
||
| it("fails when the server returns an empty response", async () => { | ||
| const { url } = setupServer((req) => { | ||
| const reqUrl = new URL(req.url); | ||
| if (reqUrl.pathname === "/api/v2/aibridge/proxy/ca-cert.pem") { | ||
| return new Response("", { status: 200 }); | ||
| } | ||
| return new Response("not found", { status: 404 }); | ||
| }); | ||
|
|
||
| const { id, instance } = await setupContainer(); | ||
|
|
||
| const exec = await execContainer(id, [ | ||
| "env", | ||
| `ACCESS_URL=${url}`, | ||
| "SESSION_TOKEN=mock-token", | ||
| "bash", | ||
| "-c", | ||
| instance.script, | ||
| ]); | ||
| expect(exec.exitCode).not.toBe(0); | ||
| expect(exec.stdout).toContain( | ||
| "AI Bridge Proxy setup failed: downloaded certificate is empty.", | ||
| ); | ||
| }); | ||
|
|
||
| it("saves the certificate to a custom path", async () => { | ||
| const { url } = setupServer((req) => { | ||
| const reqUrl = new URL(req.url); | ||
| if (reqUrl.pathname === "/api/v2/aibridge/proxy/ca-cert.pem") { | ||
| return new Response(FAKE_CERT, { | ||
| status: 200, | ||
| headers: { "Content-Type": "application/x-pem-file" }, | ||
| }); | ||
| } | ||
| return new Response("not found", { status: 404 }); | ||
| }); | ||
|
|
||
| // Pass a custom cert_path to terraform apply so the script uses it. | ||
| const { id, instance } = await setupContainer({ | ||
| cert_path: "/tmp/custom/certs/proxy-ca.pem", | ||
| }); | ||
|
|
||
| const exec = await execContainer(id, [ | ||
| "env", | ||
| `ACCESS_URL=${url}`, | ||
| "SESSION_TOKEN=mock-token", | ||
| "bash", | ||
| "-c", | ||
| instance.script, | ||
| ]); | ||
| expect(exec.exitCode).toBe(0); | ||
| expect(exec.stdout).toContain( | ||
| "AI Bridge Proxy CA certificate saved to /tmp/custom/certs/proxy-ca.pem", | ||
| ); | ||
|
|
||
| const certContent = await execContainer(id, [ | ||
| "cat", | ||
| "/tmp/custom/certs/proxy-ca.pem", | ||
| ]); | ||
| expect(certContent.stdout).toContain("BEGIN CERTIFICATE"); | ||
| }); | ||
|
|
||
| it("does not create global proxy env vars via coder_env", async () => { | ||
| const state = await runTerraformApply(import.meta.dir, { | ||
| agent_id: "foo", | ||
| proxy_url: "https://aiproxy.example.com", | ||
| }); | ||
|
|
||
| // Proxy env vars should NOT be set globally via coder_env. | ||
| // They are intended to be scoped to specific tool processes. | ||
| const proxyEnvVarNames = [ | ||
| "HTTP_PROXY", | ||
| "HTTPS_PROXY", | ||
| "NODE_EXTRA_CA_CERTS", | ||
| "SSL_CERT_FILE", | ||
| "REQUESTS_CA_BUNDLE", | ||
| "CURL_CA_BUNDLE", | ||
| ]; | ||
| const proxyEnvVars = state.resources.filter( | ||
| (r) => | ||
| r.type === "coder_env" && | ||
| r.instances.some((i) => | ||
| proxyEnvVarNames.includes(i.attributes.name as string), | ||
| ), | ||
| ); | ||
| expect(proxyEnvVars.length).toBe(0); | ||
| }); | ||
| }); |
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.