diff --git a/src/agent/Agent.zig b/src/agent/Agent.zig index e6c6e86d4..09fa7efd5 100644 --- a/src/agent/Agent.zig +++ b/src/agent/Agent.zig @@ -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(); diff --git a/src/agent/settings.zig b/src/agent/settings.zig index 8b94ea339..9e5f86c2a 100644 --- a/src/agent/settings.zig +++ b/src/agent/settings.zig @@ -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 }, }; }