From 819aee3c07c229e32e0cdafd4cf0debc3f431dcc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adri=C3=A0=20Arrufat?= Date: Tue, 28 Jul 2026 10:28:18 +0200 Subject: [PATCH] agent: support login cancellation and provider completion` --- src/agent/Agent.zig | 28 ++++++++++++++++++++++++++-- src/agent/auth/auth.zig | 7 ++++--- src/agent/auth/codex.zig | 13 +++++++++++-- 3 files changed, 41 insertions(+), 7 deletions(-) diff --git a/src/agent/Agent.zig b/src/agent/Agent.zig index 63e70d09a..13201c207 100644 --- a/src/agent/Agent.zig +++ b/src/agent/Agent.zig @@ -957,7 +957,13 @@ fn handleProvider(self: *Agent, _: std.mem.Allocator, rest: []const u8) void { // login (device-code flow) when there's none yet. if (auth.descriptorFor(provider)) |desc| { var owned = (auth.sessionFor(self.allocator, provider) catch null) orelse - (auth.login(self.allocator, desc) catch |err| { + (auth.login(self.allocator, desc, &self.cancel_requested) catch |err| { + if (err == error.LoginCancelled) { + // Undo `requestCancel`'s side effects before the next turn. + self.resetAfterCancel(self.conversation.messages.items.len); + self.terminal.printInfo("{s} login cancelled", .{desc.label}); + return; + } self.terminal.printError("{s} login failed: {s}", .{ desc.label, @errorName(err) }); return; }); @@ -1978,11 +1984,29 @@ fn completionProviders(context: *anyopaque, arena: std.mem.Allocator) []const [] }; if (reachable[i]) extra += 1; } - const names = arena.alloc([]const u8, self.available_providers.len + 1 + extra) catch return &.{}; + // Subscription providers complete even without a stored token — selecting + // one is what starts the login. + var subs: [auth.registry.len][]const u8 = undefined; + var n_subs: usize = 0; + for (auth.registry) |desc| { + const name = @tagName(desc.provider); + const detected = for (self.available_providers) |p| { + if (std.mem.eql(u8, p, name)) break true; + } else false; + if (!detected) { + subs[n_subs] = name; + n_subs += 1; + } + } + const names = arena.alloc([]const u8, self.available_providers.len + n_subs + 1 + extra) catch return &.{}; for (self.available_providers, 0..) |p, i| { names[i] = arena.dupe(u8, p) catch return &.{}; } var n = self.available_providers.len; + for (subs[0..n_subs]) |name| { + names[n] = name; + n += 1; + } for (local_providers, reachable) |tag, r| if (r) { names[n] = @tagName(tag); n += 1; diff --git a/src/agent/auth/auth.zig b/src/agent/auth/auth.zig index 0db9b28d6..52ad92152 100644 --- a/src/agent/auth/auth.zig +++ b/src/agent/auth/auth.zig @@ -84,7 +84,8 @@ pub const Descriptor = struct { device_token_url: []const u8, verify_url: []const u8, /// Interactive login (device-code flow); returns freshly-minted tokens. - loginFn: *const fn (std.mem.Allocator, *const Descriptor) anyerror!TokenSet, + /// A set `cancel` flag aborts the flow with `error.LoginCancelled`. + loginFn: *const fn (std.mem.Allocator, *const Descriptor, cancel: ?*const std.atomic.Value(bool)) anyerror!TokenSet, /// Exchange a refresh token for a new `TokenSet` (with re-derived account id). refreshFn: *const fn (std.mem.Allocator, *const Descriptor, []const u8) anyerror!TokenSet, }; @@ -233,8 +234,8 @@ pub fn sessionFor(allocator: std.mem.Allocator, provider: Config.AiProvider) !?S } /// Run the interactive login and persist the result, returning a live session. -pub fn login(allocator: std.mem.Allocator, desc: *const Descriptor) !Session { - const tokens = try desc.loginFn(allocator, desc); +pub fn login(allocator: std.mem.Allocator, desc: *const Descriptor, cancel: ?*const std.atomic.Value(bool)) !Session { + const tokens = try desc.loginFn(allocator, desc, cancel); storeSave(desc.id, tokens) catch {}; return .{ .allocator = allocator, .descriptor = desc, .tokens = tokens }; } diff --git a/src/agent/auth/codex.zig b/src/agent/auth/codex.zig index d4e0970aa..e09f0ac3c 100644 --- a/src/agent/auth/codex.zig +++ b/src/agent/auth/codex.zig @@ -41,6 +41,7 @@ const device_redirect_uri = issuer ++ "/deviceauth/callback"; const scope = "openid profile email offline_access"; const user_agent = "lightpanda"; const poll_margin_ms: u64 = 3000; +const cancel_slice_ms: u64 = 200; pub const descriptor: auth.Descriptor = .{ .provider = .codex, @@ -156,7 +157,7 @@ const DeviceToken = struct { code_verifier: []const u8, }; -fn deviceLogin(allocator: std.mem.Allocator, desc: *const auth.Descriptor) !auth.TokenSet { +fn deviceLogin(allocator: std.mem.Allocator, desc: *const auth.Descriptor, cancel: ?*const std.atomic.Value(bool)) !auth.TokenSet { _ = desc; var arena: std.heap.ArenaAllocator = .init(allocator); defer arena.deinit(); @@ -174,7 +175,15 @@ fn deviceLogin(allocator: std.mem.Allocator, desc: *const auth.Descriptor) !auth const poll_body = try std.fmt.allocPrint(a, "{{\"device_auth_id\":\"{s}\",\"user_code\":\"{s}\"}}", .{ dc.device_auth_id, dc.user_code }); const dt: DeviceToken = while (true) { - lp.io.sleep(.fromMilliseconds(@intCast(interval_ms + poll_margin_ms)), .awake) catch {}; + // The REPL's Ctrl-C only sets the cancel flag (it never kills the + // process), so the wait must poll it. + var remaining_ms: u64 = interval_ms + poll_margin_ms; + while (remaining_ms > 0) { + if (cancel) |flag| if (flag.load(.acquire)) return error.LoginCancelled; + const slice_ms = @min(remaining_ms, cancel_slice_ms); + lp.io.sleep(.fromMilliseconds(@intCast(slice_ms)), .awake) catch {}; + remaining_ms -= slice_ms; + } const res = try post(a, device_token_url, "application/json", poll_body); switch (res.status) { .ok => break try std.json.parseFromSliceLeaky(DeviceToken, a, res.body, .{ .ignore_unknown_fields = true }),