mirror of
https://github.com/lightpanda-io/browser.git
synced 2026-07-30 17:25:58 -04:00
agent: introduce KeyOwnership union for credentials
Consolidates owned_key, session, and unowned env key management into a tagged union KeyOwnership to simplify resource lifetime handling.
This commit is contained in:
@@ -135,12 +135,9 @@ const synthesis_prompt =
|
||||
allocator: std.mem.Allocator,
|
||||
ai_client: ?zenai.provider.Client,
|
||||
model_credentials: ?Credentials,
|
||||
/// Allocated credentials key (Vertex gcloud token) — other keys are unowned
|
||||
/// env pointers. The AI client references it: free only after client deinit.
|
||||
owned_key: ?[:0]const u8,
|
||||
/// Active subscription (bearer) auth. Owns `model_credentials.key`, so it must
|
||||
/// outlive the AI client; refreshed between turns by `refreshAuthIfNeeded`.
|
||||
auth_session: ?auth.Session,
|
||||
/// Ownership of `model_credentials.key`. The AI client borrows the key, so
|
||||
/// deinit only after the client is gone.
|
||||
key_ownership: settings.KeyOwnership,
|
||||
/// App data dir (owned by `App`, which outlives the agent), used to cache the
|
||||
/// models.dev catalog. Null when no data dir is available.
|
||||
app_dir: ?[]const u8,
|
||||
@@ -272,8 +269,7 @@ pub fn init(allocator: std.mem.Allocator, app: *App, opts: Config.Agent) !*Agent
|
||||
|
||||
var resolved: ?settings.ResolvedProvider = if (resolve) try settings.resolveCredentials(allocator, opts, remembered, will_repl) else null;
|
||||
// Before the ai_client errdefer, so on unwind the client goes first.
|
||||
errdefer if (resolved) |r| if (r.key_owned) allocator.free(r.credentials.key);
|
||||
errdefer if (resolved) |*r| if (r.session) |*s| s.deinit();
|
||||
errdefer if (resolved) |*r| r.ownership.deinit(allocator);
|
||||
|
||||
if (will_repl and !banner_before and resolved != null) welcome.print(resolve);
|
||||
const llm: ?Credentials = if (resolved) |r| r.credentials else null;
|
||||
@@ -330,8 +326,7 @@ pub fn init(allocator: std.mem.Allocator, app: *App, opts: Config.Agent) !*Agent
|
||||
.allocator = allocator,
|
||||
.ai_client = null,
|
||||
.model_credentials = llm,
|
||||
.owned_key = if (resolved) |r| (if (r.key_owned) r.credentials.key else null) else null,
|
||||
.auth_session = if (resolved) |r| r.session else null,
|
||||
.key_ownership = if (resolved) |r| r.ownership else .env,
|
||||
.app_dir = app.app_dir_path,
|
||||
.no_llm_persisted = remembered_no_llm,
|
||||
.model_base_url = opts.base_url,
|
||||
@@ -365,7 +360,7 @@ pub fn init(allocator: std.mem.Allocator, app: *App, opts: Config.Agent) !*Agent
|
||||
|
||||
try self.startSession();
|
||||
|
||||
self.ai_client = if (llm) |l| try zenai.provider.Client.init(lp.io, allocator, l, .{ .base_url = opts.base_url, .retry_policy = .long_running, .bill_to = hfBillTo(l.provider), .environ = lp.environ(), .account_id = if (self.auth_session) |s| s.tokens.account_id else null }) else null;
|
||||
self.ai_client = if (llm) |l| try zenai.provider.Client.init(lp.io, allocator, l, .{ .base_url = opts.base_url, .retry_policy = .long_running, .bill_to = hfBillTo(l.provider), .environ = lp.environ(), .account_id = self.key_ownership.accountId() }) else null;
|
||||
errdefer if (self.ai_client) |c| c.deinit(allocator);
|
||||
if (self.ai_client) |c| c.setInterrupt(&self.http_interrupt);
|
||||
|
||||
@@ -394,8 +389,7 @@ pub fn deinit(self: *Agent) void {
|
||||
self.browser.deinit();
|
||||
self.notification.deinit();
|
||||
if (self.ai_client) |ai_client| ai_client.deinit(self.allocator);
|
||||
if (self.auth_session) |*s| s.deinit();
|
||||
if (self.owned_key) |k| self.allocator.free(k);
|
||||
self.key_ownership.deinit(self.allocator);
|
||||
self.allocator.free(self.model);
|
||||
for (self.available_providers) |p| self.allocator.free(p);
|
||||
self.allocator.free(self.available_providers);
|
||||
@@ -947,7 +941,7 @@ fn handleProvider(self: *Agent, _: std.mem.Allocator, rest: []const u8) void {
|
||||
self.terminal.printError("could not obtain a Vertex access token: {s} (details above)", .{@errorName(err)});
|
||||
return;
|
||||
};
|
||||
self.setProvider(.{ .provider = .vertex, .key = token }, token, null) catch |err| {
|
||||
self.setProvider(.{ .provider = .vertex, .key = token }, .{ .owned = token }) catch |err| {
|
||||
self.allocator.free(token);
|
||||
self.terminal.printError("failed to set provider: {s}", .{@errorName(err)});
|
||||
};
|
||||
@@ -967,7 +961,7 @@ fn handleProvider(self: *Agent, _: std.mem.Allocator, rest: []const u8) void {
|
||||
self.terminal.printError("{s} login failed: {s}", .{ desc.label, @errorName(err) });
|
||||
return;
|
||||
});
|
||||
self.setProvider(.{ .provider = provider, .key = owned.tokens.access_token }, null, owned) catch |err| {
|
||||
self.setProvider(.{ .provider = provider, .key = owned.tokens.access_token }, .{ .session = owned }) catch |err| {
|
||||
owned.deinit();
|
||||
self.terminal.printError("failed to set provider: {s}", .{@errorName(err)});
|
||||
};
|
||||
@@ -990,7 +984,7 @@ fn handleProvider(self: *Agent, _: std.mem.Allocator, rest: []const u8) void {
|
||||
self.terminal.printError("no llama.cpp server with a loaded model at {s}", .{self.model_base_url orelse zenai.provider.llama_cpp_default_base_url});
|
||||
return;
|
||||
}
|
||||
self.setProvider(.{ .provider = provider, .key = key }, null, null) catch |err| {
|
||||
self.setProvider(.{ .provider = provider, .key = key }, .env) catch |err| {
|
||||
self.terminal.printError("failed to set provider: {s}", .{@errorName(err)});
|
||||
};
|
||||
}
|
||||
@@ -1000,8 +994,7 @@ fn handleProvider(self: *Agent, _: std.mem.Allocator, rest: []const u8) void {
|
||||
fn disableProvider(self: *Agent) void {
|
||||
if (self.ai_client) |client| client.deinit(self.allocator);
|
||||
self.ai_client = null;
|
||||
if (self.owned_key) |k| self.allocator.free(k);
|
||||
self.owned_key = null;
|
||||
self.key_ownership.deinit(self.allocator);
|
||||
self.model_credentials = null;
|
||||
self.model_completions = null;
|
||||
self.no_llm_persisted = true;
|
||||
@@ -1016,22 +1009,19 @@ fn hfBillTo(provider: Config.AiProvider) ?[]const u8 {
|
||||
return std.mem.span(v);
|
||||
}
|
||||
|
||||
/// `owned_key` transfers ownership of an allocated `credentials.key` (Vertex
|
||||
/// gcloud token) on success; on error the caller still owns it. `session`
|
||||
/// likewise transfers a subscription session that owns `credentials.key`; the
|
||||
/// previous session is freed only after the old client is gone.
|
||||
fn setProvider(self: *Agent, credentials: Credentials, owned_key: ?[:0]const u8, session: ?auth.Session) !void {
|
||||
const new_client = try zenai.provider.Client.init(lp.io, self.allocator, credentials, .{ .base_url = self.model_base_url, .retry_policy = .long_running, .bill_to = hfBillTo(credentials.provider), .environ = lp.environ(), .account_id = if (session) |s| s.tokens.account_id else null });
|
||||
/// `ownership` transfers whatever owns `credentials.key` on success; on error
|
||||
/// the caller still owns it. The previous credential is released only after
|
||||
/// the old client is gone.
|
||||
fn setProvider(self: *Agent, credentials: Credentials, ownership: settings.KeyOwnership) !void {
|
||||
const new_client = try zenai.provider.Client.init(lp.io, self.allocator, credentials, .{ .base_url = self.model_base_url, .retry_policy = .long_running, .bill_to = hfBillTo(credentials.provider), .environ = lp.environ(), .account_id = ownership.accountId() });
|
||||
errdefer new_client.deinit(self.allocator);
|
||||
|
||||
// A same-provider re-select (vertex token refresh) must not reset the model.
|
||||
const same_provider = if (self.model_credentials) |c| c.provider == credentials.provider else false;
|
||||
const new_model = try self.allocator.dupe(u8, if (same_provider) self.model else zenai.provider.defaultModel(credentials.provider));
|
||||
if (self.ai_client) |client| client.deinit(self.allocator);
|
||||
if (self.owned_key) |k| self.allocator.free(k);
|
||||
if (self.auth_session) |*s| s.deinit();
|
||||
self.owned_key = owned_key;
|
||||
self.auth_session = session;
|
||||
self.key_ownership.deinit(self.allocator);
|
||||
self.key_ownership = ownership;
|
||||
new_client.setInterrupt(&self.http_interrupt);
|
||||
self.ai_client = new_client;
|
||||
self.model_credentials = credentials;
|
||||
@@ -1059,16 +1049,16 @@ fn setProvider(self: *Agent, credentials: Credentials, owned_key: ?[:0]const u8,
|
||||
/// Runs between turns on the single agent thread, never during a request, so
|
||||
/// the displaced token is never dereferenced after it's freed here.
|
||||
fn refreshAuthIfNeeded(self: *Agent) void {
|
||||
if (self.auth_session) |*session| {
|
||||
const displaced = session.ensureFresh() catch |err| {
|
||||
self.terminal.printError("could not refresh the subscription token: {s}", .{@errorName(err)});
|
||||
return;
|
||||
};
|
||||
if (displaced) |old| {
|
||||
if (self.ai_client) |client| client.setApiKey(session.tokens.access_token);
|
||||
if (self.model_credentials) |*c| c.key = session.tokens.access_token;
|
||||
old.deinit(session.allocator);
|
||||
}
|
||||
if (self.key_ownership != .session) return;
|
||||
const session = &self.key_ownership.session;
|
||||
const displaced = session.ensureFresh() catch |err| {
|
||||
self.terminal.printError("could not refresh the subscription token: {s}", .{@errorName(err)});
|
||||
return;
|
||||
};
|
||||
if (displaced) |old| {
|
||||
if (self.ai_client) |client| client.setApiKey(session.tokens.access_token);
|
||||
if (self.model_credentials) |*c| c.key = session.tokens.access_token;
|
||||
old.deinit(session.allocator);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1633,7 +1623,7 @@ fn formatApiError(self: *Agent, client: zenai.provider.Client, err: anyerror) []
|
||||
const e = client.lastError();
|
||||
const status = e.status orelse return @errorName(err);
|
||||
const hint = if (status == 401 and client == .vertex)
|
||||
if (self.owned_key != null)
|
||||
if (self.key_ownership == .owned)
|
||||
" (Vertex token may have expired; run /provider vertex to refresh)"
|
||||
else
|
||||
" (Vertex express mode needs an express API key — a Gemini Developer key won't work)"
|
||||
@@ -1934,8 +1924,7 @@ pub fn listModels(allocator: std.mem.Allocator, opts: Config.Agent) !void {
|
||||
}
|
||||
var resolved = (try settings.resolveCredentials(allocator, opts, null, false)) orelse return error.MissingProvider;
|
||||
const llm = resolved.credentials;
|
||||
defer if (resolved.key_owned) allocator.free(llm.key);
|
||||
defer if (resolved.session) |*s| s.deinit();
|
||||
defer resolved.ownership.deinit(allocator);
|
||||
|
||||
var arena: std.heap.ArenaAllocator = .init(allocator);
|
||||
defer arena.deinit();
|
||||
|
||||
@@ -39,13 +39,35 @@ pub const api_keys_hint = "ANTHROPIC_API_KEY, OPENAI_API_KEY, GOOGLE_API_KEY, HF
|
||||
pub const ResolvedProvider = struct {
|
||||
credentials: Credentials,
|
||||
source: enum { flag, remembered, detected, picked },
|
||||
/// Key allocated (Vertex gcloud token) rather than an env pointer; the
|
||||
/// caller frees it, only after the client that references it is gone.
|
||||
key_owned: bool = false,
|
||||
/// Set for subscription (bearer) auth. Owns `credentials.key`, so it must
|
||||
/// outlive the AI client; the caller takes ownership. When set, `key_owned`
|
||||
/// stays false — the session frees the key, not the caller directly.
|
||||
session: ?auth.Session = null,
|
||||
ownership: KeyOwnership = .env,
|
||||
};
|
||||
|
||||
/// How the active `Credentials.key` is owned. `deinit` releases whatever the
|
||||
/// active variant holds — call it only after the AI client that borrows the
|
||||
/// key is gone.
|
||||
pub const KeyOwnership = union(enum) {
|
||||
/// Unowned env-var pointer.
|
||||
env,
|
||||
/// Allocated key (Vertex gcloud token); aliases the active `Credentials.key`.
|
||||
owned: [:0]const u8,
|
||||
/// Subscription (bearer) auth owning the key; refreshed between turns.
|
||||
session: auth.Session,
|
||||
|
||||
pub fn deinit(self: *KeyOwnership, allocator: std.mem.Allocator) void {
|
||||
switch (self.*) {
|
||||
.env => {},
|
||||
.owned => |k| allocator.free(k),
|
||||
.session => |*s| s.deinit(),
|
||||
}
|
||||
self.* = .env;
|
||||
}
|
||||
|
||||
pub fn accountId(self: *const KeyOwnership) ?[]const u8 {
|
||||
return switch (self.*) {
|
||||
.session => |*s| s.tokens.account_id,
|
||||
else => null,
|
||||
};
|
||||
}
|
||||
};
|
||||
|
||||
/// Probe a keyless local provider (Ollama, llama.cpp): its env key is a
|
||||
@@ -120,7 +142,7 @@ pub fn resolveCredentials(allocator: std.mem.Allocator, opts: Config.Agent, reme
|
||||
if (opts.provider) |p| {
|
||||
if (p == .vertex and vertexProjectMode()) {
|
||||
const token = try gcloudAccessToken(allocator);
|
||||
return .{ .credentials = .{ .provider = p, .key = token }, .source = .flag, .key_owned = true };
|
||||
return .{ .credentials = .{ .provider = p, .key = token }, .source = .flag, .ownership = .{ .owned = token } };
|
||||
}
|
||||
// A subscription takes priority over an API key; null = not a
|
||||
// subscription provider (or none available), so fall through.
|
||||
@@ -146,7 +168,7 @@ pub fn resolveCredentials(allocator: std.mem.Allocator, opts: Config.Agent, reme
|
||||
if (p == .vertex and vertexProjectMode()) {
|
||||
// On failure the reason is already printed; fall through to detection.
|
||||
if (gcloudAccessToken(allocator)) |token| {
|
||||
return .{ .credentials = .{ .provider = p, .key = token }, .source = .remembered, .key_owned = true };
|
||||
return .{ .credentials = .{ .provider = p, .key = token }, .source = .remembered, .ownership = .{ .owned = token } };
|
||||
} else |_| {}
|
||||
} else {
|
||||
// Subscription takes priority over an API key; both fall through to
|
||||
@@ -195,7 +217,7 @@ pub fn resolveCredentials(allocator: std.mem.Allocator, opts: Config.Agent, reme
|
||||
fn finishResolved(allocator: std.mem.Allocator, credentials: Credentials, source: @FieldType(ResolvedProvider, "source")) !ResolvedProvider {
|
||||
if (credentials.provider == .vertex and vertexProjectMode()) {
|
||||
const token = try gcloudAccessToken(allocator);
|
||||
return .{ .credentials = .{ .provider = .vertex, .key = token }, .source = source, .key_owned = true };
|
||||
return .{ .credentials = .{ .provider = .vertex, .key = token }, .source = source, .ownership = .{ .owned = token } };
|
||||
}
|
||||
if (auth.descriptorFor(credentials.provider) != null and credentials.key.len == 0) {
|
||||
if (try subscriptionResolved(allocator, credentials.provider, source)) |resolved| return resolved;
|
||||
@@ -204,10 +226,9 @@ fn finishResolved(allocator: std.mem.Allocator, credentials: Credentials, source
|
||||
return .{ .credentials = credentials, .source = source };
|
||||
}
|
||||
|
||||
/// Load a stored subscription session and wrap it as a resolved credential. The
|
||||
/// returned `session` owns `credentials.key`; the caller must keep it alive as
|
||||
/// long as the AI client and free it with `session.deinit`. Null when the user
|
||||
/// hasn't logged in.
|
||||
/// Load a stored subscription session and wrap it as a resolved credential.
|
||||
/// The session owns `credentials.key`; the caller takes ownership via
|
||||
/// `ownership.deinit`. Null when the user hasn't logged in.
|
||||
fn subscriptionResolved(allocator: std.mem.Allocator, provider: Config.AiProvider, source: @FieldType(ResolvedProvider, "source")) !?ResolvedProvider {
|
||||
const session = (try auth.sessionFor(allocator, provider)) orelse return null;
|
||||
// Name the credential in effect — a set-but-ignored API key would otherwise
|
||||
@@ -220,7 +241,7 @@ fn subscriptionResolved(allocator: std.mem.Allocator, provider: Config.AiProvide
|
||||
return .{
|
||||
.credentials = .{ .provider = provider, .key = session.tokens.access_token },
|
||||
.source = source,
|
||||
.session = session,
|
||||
.ownership = .{ .session = session },
|
||||
};
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user