Skip to content

Forward Codex compute residency headers #288

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.

Sign up for GitHub

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
alphastorm wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Forward Codex compute residency headers #288

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
137 changes: 127 additions & 10 deletions src/core/auth/chatgpt_oauth.zig
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -701,7 +701,20 @@ fn sessionFromToken(alloc: Allocator, token: *TokenSet, now_ms: i64) !chatgpt_se
return session;
}

pub fn extractAccountId(alloc: Allocator, token: []const u8) ![]u8 {
const max_codex_residency_bytes: usize = 64;

pub const RequestIdentity = struct {
account_id: []u8,
residency: ?[]u8,

pub fn deinit(self: *RequestIdentity, alloc: Allocator) void {
alloc.free(self.account_id);
if (self.residency) |residency| alloc.free(residency);
self.* = undefined;
}
};

pub fn extractRequestIdentity(alloc: Allocator, token: []const u8) !RequestIdentity {
var parts = std.mem.splitScalar(u8, token, '.');
_ = parts.next() orelse return error.InvalidChatGptAccessToken;
const payload = parts.next() orelse return error.InvalidChatGptAccessToken;
Expand All @@ -721,8 +734,36 @@ pub fn extractAccountId(alloc: Allocator, token: []const u8) ![]u8 {
if (parsed.value != .object) return error.InvalidChatGptAccessToken;
const claim = parsed.value.object.get(jwt_auth_claim) orelse return error.InvalidChatGptAccessToken;
if (claim != .object) return error.InvalidChatGptAccessToken;
return dupeRequiredString(alloc, claim.object, "chatgpt_account_id") catch

const account_id = dupeRequiredString(alloc, claim.object, "chatgpt_account_id") catch
return error.InvalidChatGptAccessToken;
errdefer alloc.free(account_id);
return .{
.account_id = account_id,
.residency = try dupeResidencyClaim(alloc, claim.object),
};
}

pub fn extractAccountId(alloc: Allocator, token: []const u8) ![]u8 {
const identity = try extractRequestIdentity(alloc, token);
if (identity.residency) |residency| alloc.free(residency);
return identity.account_id;
}

fn dupeResidencyClaim(alloc: Allocator, claim: std.json.ObjectMap) !?[]u8 {
const value = claim.get("chatgpt_compute_residency") orelse return null;
if (value != .string) return null;
const residency = std.mem.trim(u8, value.string, " \t\r\n");
if (std.mem.eql(u8, residency, "no_constraint") or !validResidencyClaim(residency)) return null;
return try alloc.dupe(u8, residency);
}

fn validResidencyClaim(residency: []const u8) bool {
if (residency.len == 0 or residency.len > max_codex_residency_bytes) return false;
for (residency) |byte| {
if (byte < 0x21 or byte > 0x7e) return false;
}
return true;
}

fn dupeRequiredString(alloc: Allocator, object: std.json.ObjectMap, key: []const u8) ![]u8 {
Expand Down Expand Up @@ -864,6 +905,20 @@ fn writeStdout(text: []const u8) !void {
try std.Io.File.stdout().writeStreamingAll(io_mod.getIo(), text);
}

fn testAccessToken(alloc: Allocator, auth_claim_json: []const u8) ![]u8 {
const payload = try std.fmt.allocPrint(
alloc,
"{{\"https://api.openai.com/auth\":{s}}}",
.{auth_claim_json},
);
defer alloc.free(payload);
const encoded_len = std.base64.url_safe_no_pad.Encoder.calcSize(payload.len);
const encoded = try alloc.alloc(u8, encoded_len);
defer alloc.free(encoded);
_ = std.base64.url_safe_no_pad.Encoder.encode(encoded, payload);
return std.fmt.allocPrint(alloc, "header.{s}.signature", .{encoded});
}

test "ChatGPT E2E OAuth endpoint overrides accept only loopback HTTP" {
try std.testing.expect(isLoopbackHttpUrl("http://127.0.0.1:1234/token"));
try std.testing.expect(isLoopbackHttpUrl("http://localhost:1234/token"));
Expand All @@ -873,21 +928,83 @@ test "ChatGPT E2E OAuth endpoint overrides accept only loopback HTTP" {

test "ChatGPT account id is extracted from the namespaced JWT claim" {
const alloc = std.testing.allocator;
const payload =
\\{"https://api.openai.com/auth":{"chatgpt_account_id":"acct_test"}}
;
const encoded_len = std.base64.url_safe_no_pad.Encoder.calcSize(payload.len);
const encoded = try alloc.alloc(u8, encoded_len);
defer alloc.free(encoded);
_ = std.base64.url_safe_no_pad.Encoder.encode(encoded, payload);
const token = try std.fmt.allocPrint(alloc, "header.{s}.signature", .{encoded});
const token = try testAccessToken(alloc, "{\"chatgpt_account_id\":\"acct_test\"}");
defer alloc.free(token);

const account_id = try extractAccountId(alloc, token);
defer alloc.free(account_id);
try std.testing.expectEqualStrings("acct_test", account_id);
}

test "ChatGPT request identity extracts bounded compute residency" {
const alloc = std.testing.allocator;
const max_residency = "r" ** max_codex_residency_bytes;
const max_claim = try std.fmt.allocPrint(
alloc,
"{{\"chatgpt_account_id\":\"acct_test\",\"chatgpt_compute_residency\":\"{s}\"}}",
.{max_residency},
);
defer alloc.free(max_claim);
const cases = [_]struct {
claim: []const u8,
expected_residency: ?[]const u8,
}{
.{
.claim = "{\"chatgpt_account_id\":\"acct_test\",\"chatgpt_data_residency\":\"eu\",\"chatgpt_compute_residency\":\"us\"}",
.expected_residency = "us",
},
.{
.claim = "{\"chatgpt_account_id\":\"acct_test\",\"chatgpt_compute_residency\":\" eu \"}",
.expected_residency = "eu",
},
.{ .claim = max_claim, .expected_residency = max_residency },
.{
.claim = "{\"chatgpt_account_id\":\"acct_test\",\"chatgpt_data_residency\":\"eu\"}",
.expected_residency = null,
},
.{ .claim = "{\"chatgpt_account_id\":\"acct_test\"}", .expected_residency = null },
};

for (cases) |case| {
const token = try testAccessToken(alloc, case.claim);
defer alloc.free(token);
var identity = try extractRequestIdentity(alloc, token);
defer identity.deinit(alloc);
try std.testing.expectEqualStrings("acct_test", identity.account_id);
if (case.expected_residency) |expected| {
try std.testing.expectEqualStrings(expected, identity.residency orelse "");
} else {
try std.testing.expect(identity.residency == null);
}
}
}

test "ChatGPT request identity omits unconstrained and unsafe compute residency" {
const alloc = std.testing.allocator;
const overlong_residency = "r" ** (max_codex_residency_bytes + 1);
const overlong_claim = try std.fmt.allocPrint(
alloc,
"{{\"chatgpt_account_id\":\"acct_test\",\"chatgpt_compute_residency\":\"{s}\"}}",
.{overlong_residency},
);
defer alloc.free(overlong_claim);
const claims = [_][]const u8{
"{\"chatgpt_account_id\":\"acct_test\",\"chatgpt_compute_residency\":\"no_constraint\"}",
"{\"chatgpt_account_id\":\"acct_test\",\"chatgpt_compute_residency\":\"\"}",
"{\"chatgpt_account_id\":\"acct_test\",\"chatgpt_compute_residency\":\"us\\r\\nx-injected: yes\"}",
"{\"chatgpt_account_id\":\"acct_test\",\"chatgpt_compute_residency\":42}",
overlong_claim,
};

for (claims) |claim| {
const token = try testAccessToken(alloc, claim);
defer alloc.free(token);
var identity = try extractRequestIdentity(alloc, token);
defer identity.deinit(alloc);
try std.testing.expect(identity.residency == null);
}
}

test "Codex refresh uses JSON and accepts omitted token rotation and lifetime" {
const State = struct {
method: ?oauth_transport.Method = null,
Expand Down
10 changes: 7 additions & 3 deletions src/gateway/openai_codex.zig
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -326,8 +326,8 @@ fn streamCompletionCore(alloc: Allocator, request: stream_provider.Request) !str
return error.CodexSubscriptionCredentialRequired;
}
try validateModel(request.model);
const account_id = try chatgpt_oauth.extractAccountId(alloc, request.api_key);
defer alloc.free(account_id);
var request_identity = try chatgpt_oauth.extractRequestIdentity(alloc, request.api_key);
defer request_identity.deinit(alloc);
const auth_header = try std.fmt.allocPrint(alloc, "Bearer {s}", .{request.api_key});
defer secret.zeroAndFree(alloc, auth_header);
const request_endpoint = if (io_mod.getenv(e2e_endpoint_env)) |override| endpoint: {
Expand All @@ -338,8 +338,12 @@ fn streamCompletionCore(alloc: Allocator, request: stream_provider.Request) !str

var extra_headers_buf: [7]std.http.Header = undefined;
var extra_count: usize = 0;
extra_headers_buf[extra_count] = .{ .name = "chatgpt-account-id", .value = account_id };
extra_headers_buf[extra_count] = .{ .name = "chatgpt-account-id", .value = request_identity.account_id };
extra_count += 1;
if (request_identity.residency) |residency| {
extra_headers_buf[extra_count] = .{ .name = "x-openai-internal-codex-residency", .value = residency };
extra_count += 1;
}
extra_headers_buf[extra_count] = .{ .name = "originator", .value = "fx" };
extra_count += 1;
extra_headers_buf[extra_count] = .{ .name = "OpenAI-Beta", .value = "responses=experimental" };
Expand Down
26 changes: 18 additions & 8 deletions src/gateway/openai_codex_models.zig
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,11 +63,11 @@ fn fetchCatalogForProvider(
}
const credential = input.access.authorizationCredential() orelse
return .{ .failure = .{ .category = .authentication, .http_status = .unauthorized } };
const account_id = chatgpt_oauth.extractAccountId(alloc, credential) catch |err| {
var request_identity = chatgpt_oauth.extractRequestIdentity(alloc, credential) catch |err| {
if (err == error.OutOfMemory) return error.OutOfMemory;
return .{ .failure = .{ .category = .authentication, .http_status = .unauthorized } };
};
defer alloc.free(account_id);
defer request_identity.deinit(alloc);
const request_url = modelsUrl(alloc) catch |err| {
if (err == error.OutOfMemory) return error.OutOfMemory;
return .{ .failure = .{ .category = .runtime } };
Expand All @@ -80,7 +80,8 @@ fn fetchCatalogForProvider(
.alloc = alloc,
.url = request_url,
.credential = credential,
.account_id = account_id,
.account_id = request_identity.account_id,
.residency = request_identity.residency,
};
var response = gateway_client.runBoundedHttpOperation(
FetchResponse,
Expand Down Expand Up @@ -135,6 +136,7 @@ const FetchOperation = struct {
url: []const u8,
credential: []const u8,
account_id: []const u8,
residency: ?[]const u8,

pub fn run(self: *@This()) !FetchResponse {
var client: std.http.Client = .{ .allocator = self.alloc, .io = io_mod.getIo() };
Expand All @@ -144,6 +146,18 @@ const FetchOperation = struct {
const body_buffer = try self.alloc.alloc(u8, max_catalog_bytes + 1);
defer secret.zeroAndFree(self.alloc, body_buffer);
var response_writer = std.Io.Writer.fixed(body_buffer);
var extra_headers_buf: [4]std.http.Header = undefined;
var extra_count: usize = 0;
extra_headers_buf[extra_count] = .{ .name = "chatgpt-account-id", .value = self.account_id };
extra_count += 1;
if (self.residency) |residency| {
extra_headers_buf[extra_count] = .{ .name = "x-openai-internal-codex-residency", .value = residency };
extra_count += 1;
}
extra_headers_buf[extra_count] = .{ .name = "originator", .value = "fx" };
extra_count += 1;
extra_headers_buf[extra_count] = .{ .name = "accept", .value = "application/json" };
extra_count += 1;
const result = client.fetch(.{
.location = .{ .url = self.url },
.method = .GET,
Expand All @@ -152,11 +166,7 @@ const FetchOperation = struct {
.user_agent = .{ .override = gateway_client.user_agent },
.accept_encoding = .omit,
},
.extra_headers = &.{
.{ .name = "chatgpt-account-id", .value = self.account_id },
.{ .name = "originator", .value = "fx" },
.{ .name = "accept", .value = "application/json" },
},
.extra_headers = extra_headers_buf[0..extra_count],
.response_writer = &response_writer,
.redirect_behavior = .unhandled,
}) catch |err| switch (err) {
Expand Down
70 changes: 63 additions & 7 deletions tests/e2e/acp.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -522,9 +522,13 @@ function writeSeededFxAuth(home: string, teamId?: string): void {
function acpChatGptAccessToken(
accountId = "acct_acp_e2e",
signature = "signature",
authClaims: Record<string, unknown> = {},
): string {
const payload = Buffer.from(JSON.stringify({
"https://api.openai.com/auth": { chatgpt_account_id: accountId },
"https://api.openai.com/auth": {
chatgpt_account_id: accountId,
...authClaims,
},
})).toString("base64url");
return `header.${payload}.${signature}`;
}
Expand Down Expand Up @@ -577,19 +581,24 @@ function codexLatestToolResult(body: string): { callId: string; output: string }
function startAcpFakeCodex(options: {
unauthorizedResponses?: number;
route?: (body: string) => string | Promise<string>;
authClaims?: Record<string, unknown>;
} = {}) {
const accessToken = acpChatGptAccessToken("acct_acp_e2e", "stale");
const refreshedAccessToken = acpChatGptAccessToken("acct_acp_e2e", "fresh");
const requests: Array<{ path: string; authorization: string | null; body: string }> = [];
const modelRequests: Array<{ path: string; authorization: string | null }> = [];
const tokenRequests: Array<{ path: string; authorization: string | null }> = [];
const accessToken = acpChatGptAccessToken("acct_acp_e2e", "stale", options.authClaims);
const refreshedAccessToken = acpChatGptAccessToken("acct_acp_e2e", "fresh", options.authClaims);
const requests: Array<{ path: string; authorization: string | null; residency: string | null; body: string }> = [];
const modelRequests: Array<{ path: string; authorization: string | null; residency: string | null }> = [];
const tokenRequests: Array<{ path: string; authorization: string | null; residency: string | null }> = [];
let unauthorizedResponses = options.unauthorizedResponses ?? 0;
const server = Bun.serve({
hostname: "127.0.0.1",
port: 0,
async fetch(request) {
const path = new URL(request.url).pathname;
const recorded = { path, authorization: request.headers.get("authorization") };
const recorded = {
path,
authorization: request.headers.get("authorization"),
residency: request.headers.get("x-openai-internal-codex-residency"),
};
if (path === "/models") {
modelRequests.push(recorded);
return Response.json({ models: [
Expand Down Expand Up @@ -7462,6 +7471,53 @@ describe("acp: model catalog authentication", () => {
},
TIMEOUT,
);

test(
"session provider changes propagate Codex residency to model discovery",
async () => {
const root = createIsolatedRoot("fx-acp-codex-residency-");
const gateway = startFakeGateway([]);
const codex = startAcpFakeCodex({
authClaims: { chatgpt_compute_residency: "us" },
});
writeSeededAcpChatGptLogin(root.home, codex.accessToken);
try {
client = await AcpClient.create({
cwd: root.workspace,
env: {
...fakeGatewayEnv(root, gateway),
FX_E2E_OPENAI_CODEX_RESPONSES_URL: codex.responsesUrl,
FX_E2E_OPENAI_CODEX_MODELS_URL: codex.modelsUrl,
FX_E2E_CHATGPT_TOKEN_URL: codex.tokenUrl,
},
});
await client.request("initialize", { protocolVersion: 1 }, 1);
await client.request("session/new", { mcpServers: [] }, 2);
await client.readLine();

const changed = await client.request("session/set_config_option", {
configId: "provider",
value: "codex",
}, 3) as any;
expect(changed.result.configOptions.find((option: any) => option.id === "provider").currentValue)
.toBe("codex");
expect(codex.modelRequests).toHaveLength(1);
expect(codex.modelRequests[0]!.residency).toBe("us");
expect(gateway.requests).toHaveLength(0);
expect(gateway.modelRequests).toHaveLength(1);
for (const request of [...gateway.requests, ...gateway.modelRequests]) {
expect(request.headers.has("x-openai-internal-codex-residency")).toBe(false);
}
expect(client.stderr).toBe("");
} finally {
await client?.close();
codex.stop();
gateway.stop();
rmSync(root.root, { recursive: true, force: true });
}
},
TIMEOUT,
);
});

describe.skipIf(!HAS_API_KEY)("acp: model-backed protocol", () => {
Expand Down
Loading