klaude-code 2.6.0__py3-none-any.whl → 2.8.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (82) hide show
  1. klaude_code/app/runtime.py +1 -1
  2. klaude_code/auth/AGENTS.md +325 -0
  3. klaude_code/auth/__init__.py +17 -1
  4. klaude_code/auth/antigravity/__init__.py +20 -0
  5. klaude_code/auth/antigravity/exceptions.py +17 -0
  6. klaude_code/auth/antigravity/oauth.py +320 -0
  7. klaude_code/auth/antigravity/pkce.py +25 -0
  8. klaude_code/auth/antigravity/token_manager.py +45 -0
  9. klaude_code/auth/base.py +4 -0
  10. klaude_code/auth/claude/oauth.py +29 -9
  11. klaude_code/auth/codex/exceptions.py +4 -0
  12. klaude_code/auth/env.py +19 -15
  13. klaude_code/cli/auth_cmd.py +54 -4
  14. klaude_code/cli/cost_cmd.py +83 -160
  15. klaude_code/cli/list_model.py +50 -0
  16. klaude_code/cli/main.py +99 -9
  17. klaude_code/config/assets/builtin_config.yaml +108 -0
  18. klaude_code/config/builtin_config.py +5 -11
  19. klaude_code/config/config.py +24 -10
  20. klaude_code/const.py +11 -1
  21. klaude_code/core/agent.py +5 -1
  22. klaude_code/core/agent_profile.py +28 -32
  23. klaude_code/core/compaction/AGENTS.md +112 -0
  24. klaude_code/core/compaction/__init__.py +11 -0
  25. klaude_code/core/compaction/compaction.py +707 -0
  26. klaude_code/core/compaction/overflow.py +30 -0
  27. klaude_code/core/compaction/prompts.py +97 -0
  28. klaude_code/core/executor.py +103 -2
  29. klaude_code/core/manager/llm_clients.py +5 -0
  30. klaude_code/core/manager/llm_clients_builder.py +14 -2
  31. klaude_code/core/prompts/prompt-antigravity.md +80 -0
  32. klaude_code/core/prompts/prompt-codex-gpt-5-2.md +335 -0
  33. klaude_code/core/reminders.py +11 -7
  34. klaude_code/core/task.py +126 -0
  35. klaude_code/core/tool/todo/todo_write_tool.py +1 -1
  36. klaude_code/core/turn.py +3 -1
  37. klaude_code/llm/antigravity/__init__.py +3 -0
  38. klaude_code/llm/antigravity/client.py +558 -0
  39. klaude_code/llm/antigravity/input.py +261 -0
  40. klaude_code/llm/registry.py +1 -0
  41. klaude_code/protocol/commands.py +0 -1
  42. klaude_code/protocol/events.py +18 -0
  43. klaude_code/protocol/llm_param.py +1 -0
  44. klaude_code/protocol/message.py +23 -1
  45. klaude_code/protocol/op.py +15 -1
  46. klaude_code/protocol/op_handler.py +5 -0
  47. klaude_code/session/session.py +36 -0
  48. klaude_code/skill/assets/create-plan/SKILL.md +6 -6
  49. klaude_code/skill/loader.py +12 -13
  50. klaude_code/skill/manager.py +3 -3
  51. klaude_code/tui/command/__init__.py +4 -4
  52. klaude_code/tui/command/compact_cmd.py +32 -0
  53. klaude_code/tui/command/copy_cmd.py +1 -1
  54. klaude_code/tui/command/fork_session_cmd.py +114 -18
  55. klaude_code/tui/command/model_picker.py +5 -1
  56. klaude_code/tui/command/thinking_cmd.py +1 -1
  57. klaude_code/tui/commands.py +6 -0
  58. klaude_code/tui/components/command_output.py +1 -1
  59. klaude_code/tui/components/rich/markdown.py +117 -1
  60. klaude_code/tui/components/rich/theme.py +18 -2
  61. klaude_code/tui/components/tools.py +39 -25
  62. klaude_code/tui/components/user_input.py +39 -28
  63. klaude_code/tui/input/AGENTS.md +44 -0
  64. klaude_code/tui/input/__init__.py +5 -2
  65. klaude_code/tui/input/completers.py +10 -14
  66. klaude_code/tui/input/drag_drop.py +146 -0
  67. klaude_code/tui/input/images.py +227 -0
  68. klaude_code/tui/input/key_bindings.py +183 -19
  69. klaude_code/tui/input/paste.py +71 -0
  70. klaude_code/tui/input/prompt_toolkit.py +32 -9
  71. klaude_code/tui/machine.py +26 -1
  72. klaude_code/tui/renderer.py +67 -4
  73. klaude_code/tui/runner.py +19 -3
  74. klaude_code/tui/terminal/image.py +103 -10
  75. klaude_code/tui/terminal/selector.py +81 -7
  76. {klaude_code-2.6.0.dist-info → klaude_code-2.8.0.dist-info}/METADATA +10 -10
  77. {klaude_code-2.6.0.dist-info → klaude_code-2.8.0.dist-info}/RECORD +79 -61
  78. klaude_code/core/prompts/prompt-codex-gpt-5-1-codex-max.md +0 -117
  79. klaude_code/tui/command/terminal_setup_cmd.py +0 -248
  80. klaude_code/tui/input/clipboard.py +0 -152
  81. {klaude_code-2.6.0.dist-info → klaude_code-2.8.0.dist-info}/WHEEL +0 -0
  82. {klaude_code-2.6.0.dist-info → klaude_code-2.8.0.dist-info}/entry_points.txt +0 -0
@@ -0,0 +1,320 @@
1
+ """OAuth PKCE flow for Antigravity authentication."""
2
+
3
+ import base64
4
+ import json
5
+ import secrets
6
+ import time
7
+ import webbrowser
8
+ from http.server import BaseHTTPRequestHandler, HTTPServer
9
+ from threading import Thread
10
+ from typing import Any, cast
11
+ from urllib.parse import parse_qs, urlencode, urlparse
12
+
13
+ import httpx
14
+
15
+ from klaude_code.auth.antigravity.exceptions import (
16
+ AntigravityNotLoggedInError,
17
+ AntigravityOAuthError,
18
+ AntigravityTokenExpiredError,
19
+ )
20
+ from klaude_code.auth.antigravity.pkce import generate_pkce
21
+ from klaude_code.auth.antigravity.token_manager import AntigravityAuthState, AntigravityTokenManager
22
+
23
+ # OAuth configuration (decoded from base64 for compatibility with pi implementation)
24
+ CLIENT_ID = base64.b64decode(
25
+ "MTA3MTAwNjA2MDU5MS10bWhzc2luMmgyMWxjcmUyMzV2dG9sb2poNGc0MDNlcC5hcHBzLmdvb2dsZXVzZXJjb250ZW50LmNvbQ=="
26
+ ).decode()
27
+ CLIENT_SECRET = base64.b64decode("R09DU1BYLUs1OEZXUjQ4NkxkTEoxbUxCOHNYQzR6NnFEQWY=").decode()
28
+ REDIRECT_URI = "http://localhost:51121/oauth-callback"
29
+ REDIRECT_PORT = 51121
30
+
31
+ # Google OAuth endpoints
32
+ AUTH_URL = "https://accounts.google.com/o/oauth2/v2/auth"
33
+ TOKEN_URL = "https://oauth2.googleapis.com/token"
34
+
35
+ # Antigravity requires additional scopes
36
+ SCOPES = [
37
+ "https://www.googleapis.com/auth/cloud-platform",
38
+ "https://www.googleapis.com/auth/userinfo.email",
39
+ "https://www.googleapis.com/auth/userinfo.profile",
40
+ "https://www.googleapis.com/auth/cclog",
41
+ "https://www.googleapis.com/auth/experimentsandconfigs",
42
+ ]
43
+
44
+ # Fallback project ID when discovery fails
45
+ DEFAULT_PROJECT_ID = "rising-fact-p41fc"
46
+
47
+ # Cloud Code Assist endpoint
48
+ CLOUDCODE_ENDPOINT = "https://cloudcode-pa.googleapis.com"
49
+
50
+
51
+ class OAuthCallbackHandler(BaseHTTPRequestHandler):
52
+ """HTTP request handler for OAuth callback."""
53
+
54
+ code: str | None = None
55
+ state: str | None = None
56
+ error: str | None = None
57
+
58
+ def log_message(self, format: str, *args: Any) -> None:
59
+ """Suppress HTTP server logs."""
60
+ pass
61
+
62
+ def do_GET(self) -> None:
63
+ """Handle GET request from OAuth callback."""
64
+ parsed = urlparse(self.path)
65
+ params = parse_qs(parsed.query)
66
+
67
+ OAuthCallbackHandler.code = params.get("code", [None])[0]
68
+ OAuthCallbackHandler.state = params.get("state", [None])[0]
69
+ OAuthCallbackHandler.error = params.get("error", [None])[0]
70
+
71
+ self.send_response(200)
72
+ self.send_header("Content-Type", "text/html")
73
+ self.end_headers()
74
+
75
+ if OAuthCallbackHandler.error:
76
+ html = f"""
77
+ <html><body style="font-family: sans-serif; text-align: center; padding: 50px;">
78
+ <h1>Authentication Failed</h1>
79
+ <p>Error: {OAuthCallbackHandler.error}</p>
80
+ <p>Please close this window and try again.</p>
81
+ </body></html>
82
+ """
83
+ else:
84
+ html = """
85
+ <html><body style="font-family: sans-serif; text-align: center; padding: 50px;">
86
+ <h1>Authentication Successful!</h1>
87
+ <p>You can close this window now.</p>
88
+ <script>setTimeout(function() { window.close(); }, 2000);</script>
89
+ </body></html>
90
+ """
91
+ self.wfile.write(html.encode())
92
+
93
+
94
+ def _discover_project(access_token: str) -> str:
95
+ """Discover or provision a project for the user."""
96
+ headers = {
97
+ "Authorization": f"Bearer {access_token}",
98
+ "Content-Type": "application/json",
99
+ "User-Agent": "google-api-nodejs-client/9.15.1",
100
+ "X-Goog-Api-Client": "google-cloud-sdk vscode_cloudshelleditor/0.1",
101
+ "Client-Metadata": json.dumps(
102
+ {
103
+ "ideType": "IDE_UNSPECIFIED",
104
+ "platform": "PLATFORM_UNSPECIFIED",
105
+ "pluginType": "GEMINI",
106
+ }
107
+ ),
108
+ }
109
+
110
+ try:
111
+ with httpx.Client() as client:
112
+ response = client.post(
113
+ f"{CLOUDCODE_ENDPOINT}/v1internal:loadCodeAssist",
114
+ headers=headers,
115
+ json={
116
+ "metadata": {
117
+ "ideType": "IDE_UNSPECIFIED",
118
+ "platform": "PLATFORM_UNSPECIFIED",
119
+ "pluginType": "GEMINI",
120
+ },
121
+ },
122
+ timeout=30,
123
+ )
124
+
125
+ if response.status_code == 200:
126
+ data: dict[str, Any] = response.json()
127
+ project = data.get("cloudaicompanionProject")
128
+ if isinstance(project, str) and project:
129
+ return project
130
+ if isinstance(project, dict):
131
+ project_dict = cast(dict[str, Any], project)
132
+ project_id = project_dict.get("id")
133
+ if isinstance(project_id, str) and project_id:
134
+ return project_id
135
+ except Exception:
136
+ pass
137
+
138
+ return DEFAULT_PROJECT_ID
139
+
140
+
141
+ def _get_user_email(access_token: str) -> str | None:
142
+ """Get user email from the access token."""
143
+ try:
144
+ with httpx.Client() as client:
145
+ response = client.get(
146
+ "https://www.googleapis.com/oauth2/v1/userinfo?alt=json",
147
+ headers={"Authorization": f"Bearer {access_token}"},
148
+ timeout=10,
149
+ )
150
+ if response.status_code == 200:
151
+ data = response.json()
152
+ return data.get("email")
153
+ except Exception:
154
+ pass
155
+ return None
156
+
157
+
158
+ class AntigravityOAuth:
159
+ """Handle OAuth PKCE flow for Antigravity authentication."""
160
+
161
+ def __init__(self, token_manager: AntigravityTokenManager | None = None):
162
+ self.token_manager = token_manager or AntigravityTokenManager()
163
+
164
+ def login(self) -> AntigravityAuthState:
165
+ """Run the complete OAuth login flow."""
166
+ verifier, challenge = generate_pkce()
167
+ state = secrets.token_urlsafe(32)
168
+
169
+ # Build authorization URL
170
+ auth_params = {
171
+ "client_id": CLIENT_ID,
172
+ "response_type": "code",
173
+ "redirect_uri": REDIRECT_URI,
174
+ "scope": " ".join(SCOPES),
175
+ "code_challenge": challenge,
176
+ "code_challenge_method": "S256",
177
+ "state": state,
178
+ "access_type": "offline",
179
+ "prompt": "consent",
180
+ }
181
+ auth_url = f"{AUTH_URL}?{urlencode(auth_params)}"
182
+
183
+ # Reset callback handler state
184
+ OAuthCallbackHandler.code = None
185
+ OAuthCallbackHandler.state = None
186
+ OAuthCallbackHandler.error = None
187
+
188
+ # Start callback server
189
+ server = HTTPServer(("localhost", REDIRECT_PORT), OAuthCallbackHandler)
190
+ server_thread = Thread(target=server.handle_request)
191
+ server_thread.start()
192
+
193
+ # Open browser for user to authenticate
194
+ webbrowser.open(auth_url)
195
+
196
+ # Wait for callback
197
+ server_thread.join(timeout=300) # 5 minute timeout
198
+ server.server_close()
199
+
200
+ # Check for errors
201
+ if OAuthCallbackHandler.error:
202
+ raise AntigravityOAuthError(f"OAuth error: {OAuthCallbackHandler.error}")
203
+
204
+ if not OAuthCallbackHandler.code:
205
+ raise AntigravityOAuthError("No authorization code received")
206
+
207
+ if OAuthCallbackHandler.state is None or OAuthCallbackHandler.state != state:
208
+ raise AntigravityOAuthError("OAuth state mismatch")
209
+
210
+ # Exchange code for tokens
211
+ auth_state = self._exchange_code(OAuthCallbackHandler.code, verifier)
212
+
213
+ # Save tokens
214
+ self.token_manager.save(auth_state)
215
+
216
+ return auth_state
217
+
218
+ def _exchange_code(self, code: str, verifier: str) -> AntigravityAuthState:
219
+ """Exchange authorization code for tokens."""
220
+ data = {
221
+ "client_id": CLIENT_ID,
222
+ "client_secret": CLIENT_SECRET,
223
+ "code": code,
224
+ "grant_type": "authorization_code",
225
+ "redirect_uri": REDIRECT_URI,
226
+ "code_verifier": verifier,
227
+ }
228
+
229
+ with httpx.Client() as client:
230
+ response = client.post(TOKEN_URL, data=data, timeout=30)
231
+
232
+ if response.status_code != 200:
233
+ raise AntigravityOAuthError(f"Token exchange failed: {response.text}")
234
+
235
+ tokens = response.json()
236
+ access_token = tokens["access_token"]
237
+ refresh_token = tokens.get("refresh_token")
238
+ expires_in = tokens.get("expires_in", 3600)
239
+
240
+ if not refresh_token:
241
+ raise AntigravityOAuthError("No refresh token received. Please try again.")
242
+
243
+ # Get user email
244
+ email = _get_user_email(access_token)
245
+
246
+ # Discover project
247
+ project_id = _discover_project(access_token)
248
+
249
+ # Calculate expiry time with 5 minute buffer
250
+ expires_at = int(time.time()) + expires_in - 300
251
+
252
+ return AntigravityAuthState(
253
+ access_token=access_token,
254
+ refresh_token=refresh_token,
255
+ expires_at=expires_at,
256
+ project_id=project_id,
257
+ email=email,
258
+ )
259
+
260
+ def refresh(self) -> AntigravityAuthState:
261
+ """Refresh the access token using refresh token."""
262
+ state = self.token_manager.get_state()
263
+ if state is None:
264
+ raise AntigravityNotLoggedInError("Not logged in to Antigravity. Run 'klaude login antigravity' first.")
265
+
266
+ data = {
267
+ "client_id": CLIENT_ID,
268
+ "client_secret": CLIENT_SECRET,
269
+ "refresh_token": state.refresh_token,
270
+ "grant_type": "refresh_token",
271
+ }
272
+
273
+ with httpx.Client() as client:
274
+ response = client.post(TOKEN_URL, data=data, timeout=30)
275
+
276
+ if response.status_code != 200:
277
+ raise AntigravityTokenExpiredError(f"Token refresh failed: {response.text}")
278
+
279
+ tokens = response.json()
280
+ access_token = tokens["access_token"]
281
+ refresh_token = tokens.get("refresh_token", state.refresh_token)
282
+ expires_in = tokens.get("expires_in", 3600)
283
+
284
+ # Calculate expiry time with 5 minute buffer
285
+ expires_at = int(time.time()) + expires_in - 300
286
+
287
+ new_state = AntigravityAuthState(
288
+ access_token=access_token,
289
+ refresh_token=refresh_token,
290
+ expires_at=expires_at,
291
+ project_id=state.project_id,
292
+ email=state.email,
293
+ )
294
+
295
+ self.token_manager.save(new_state)
296
+ return new_state
297
+
298
+ def ensure_valid_token(self) -> tuple[str, str]:
299
+ """Ensure we have a valid access token, refreshing if needed.
300
+
301
+ Returns:
302
+ Tuple of (access_token, project_id).
303
+ """
304
+ state = self.token_manager.get_state()
305
+ if state is None:
306
+ raise AntigravityNotLoggedInError("Not logged in to Antigravity. Run 'klaude login antigravity' first.")
307
+
308
+ if state.is_expired():
309
+ state = self.refresh()
310
+
311
+ return state.access_token, state.project_id
312
+
313
+ def get_api_key_json(self) -> str:
314
+ """Get API key as JSON string for LLM client.
315
+
316
+ Returns:
317
+ JSON string with token and projectId.
318
+ """
319
+ access_token, project_id = self.ensure_valid_token()
320
+ return json.dumps({"token": access_token, "projectId": project_id})
@@ -0,0 +1,25 @@
1
+ """PKCE utilities for Antigravity OAuth."""
2
+
3
+ import base64
4
+ import hashlib
5
+ import secrets
6
+
7
+
8
+ def base64url_encode(data: bytes) -> str:
9
+ """Encode bytes as base64url string without padding."""
10
+ return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")
11
+
12
+
13
+ def generate_pkce() -> tuple[str, str]:
14
+ """Generate PKCE code verifier and challenge.
15
+
16
+ Returns:
17
+ Tuple of (verifier, challenge).
18
+ """
19
+ verifier_bytes = secrets.token_bytes(32)
20
+ verifier = base64url_encode(verifier_bytes)
21
+
22
+ challenge_hash = hashlib.sha256(verifier.encode("ascii")).digest()
23
+ challenge = base64url_encode(challenge_hash)
24
+
25
+ return verifier, challenge
@@ -0,0 +1,45 @@
1
+ """Token storage and management for Antigravity authentication."""
2
+
3
+ from pathlib import Path
4
+ from typing import Any
5
+
6
+ from klaude_code.auth.base import BaseAuthState, BaseTokenManager
7
+
8
+
9
+ class AntigravityAuthState(BaseAuthState):
10
+ """Stored authentication state for Antigravity."""
11
+
12
+ project_id: str
13
+ email: str | None = None
14
+
15
+
16
+ class AntigravityTokenManager(BaseTokenManager[AntigravityAuthState]):
17
+ """Manage Antigravity OAuth tokens."""
18
+
19
+ def __init__(self, auth_file: Path | None = None):
20
+ super().__init__(auth_file)
21
+
22
+ @property
23
+ def storage_key(self) -> str:
24
+ return "antigravity"
25
+
26
+ def _create_state(self, data: dict[str, Any]) -> AntigravityAuthState:
27
+ return AntigravityAuthState.model_validate(data)
28
+
29
+ def get_access_token(self) -> str:
30
+ """Get access token, raising if not logged in."""
31
+ state = self.get_state()
32
+ if state is None:
33
+ from klaude_code.auth.antigravity.exceptions import AntigravityNotLoggedInError
34
+
35
+ raise AntigravityNotLoggedInError("Not logged in to Antigravity. Run 'klaude login antigravity' first.")
36
+ return state.access_token
37
+
38
+ def get_project_id(self) -> str:
39
+ """Get project ID, raising if not logged in."""
40
+ state = self.get_state()
41
+ if state is None:
42
+ from klaude_code.auth.antigravity.exceptions import AntigravityNotLoggedInError
43
+
44
+ raise AntigravityNotLoggedInError("Not logged in to Antigravity. Run 'klaude login antigravity' first.")
45
+ return state.project_id
klaude_code/auth/base.py CHANGED
@@ -95,3 +95,7 @@ class BaseTokenManager[T: BaseAuthState](ABC):
95
95
  if self._state is None:
96
96
  self._state = self.load()
97
97
  return self._state
98
+
99
+ def clear_cached_state(self) -> None:
100
+ """Clear in-memory cached state to force reload from file on next access."""
101
+ self._state = None
@@ -125,25 +125,45 @@ class ClaudeOAuth:
125
125
  expires_at=int(time.time()) + int(expires_in),
126
126
  )
127
127
 
128
- def refresh(self) -> ClaudeAuthState:
129
- """Refresh the access token using refresh token."""
130
- state = self.token_manager.get_state()
131
- if state is None:
132
- raise ClaudeNotLoggedInError("Not logged in to Claude. Run 'klaude login claude' first.")
133
-
128
+ def _do_refresh_request(self, refresh_token: str) -> httpx.Response:
129
+ """Send token refresh request to OAuth server."""
134
130
  payload = {
135
131
  "grant_type": "refresh_token",
136
132
  "client_id": CLIENT_ID,
137
- "refresh_token": state.refresh_token,
133
+ "refresh_token": refresh_token,
138
134
  }
139
-
140
135
  with httpx.Client() as client:
141
- response = client.post(
136
+ return client.post(
142
137
  TOKEN_URL,
143
138
  json=payload,
144
139
  headers={"Content-Type": "application/json"},
145
140
  )
146
141
 
142
+ def refresh(self) -> ClaudeAuthState:
143
+ """Refresh the access token using refresh token.
144
+
145
+ Handles concurrent refresh race conditions by retrying with freshly loaded token
146
+ if the first attempt fails with invalid_grant error.
147
+ """
148
+ state = self.token_manager.get_state()
149
+ if state is None:
150
+ raise ClaudeNotLoggedInError("Not logged in to Claude. Run 'klaude login claude' first.")
151
+
152
+ response = self._do_refresh_request(state.refresh_token)
153
+
154
+ # Handle race condition: another process may have refreshed the token already
155
+ if response.status_code != 200 and "invalid_grant" in response.text:
156
+ # Reload token from file (another process may have updated it)
157
+ self.token_manager.clear_cached_state()
158
+ fresh_state = self.token_manager.load()
159
+ if fresh_state and fresh_state.refresh_token != state.refresh_token:
160
+ # Token was updated by another process
161
+ if not fresh_state.is_expired():
162
+ # New token is still valid, use it directly
163
+ return fresh_state
164
+ # New token expired, try refreshing with the new refresh_token
165
+ response = self._do_refresh_request(fresh_state.refresh_token)
166
+
147
167
  if response.status_code != 200:
148
168
  raise ClaudeAuthError(f"Token refresh failed: {response.text}")
149
169
 
@@ -15,3 +15,7 @@ class CodexTokenExpiredError(CodexAuthError):
15
15
 
16
16
  class CodexOAuthError(CodexAuthError):
17
17
  """OAuth flow failed."""
18
+
19
+
20
+ class CodexUnsupportedModelError(CodexAuthError):
21
+ """Model is not supported by codex_oauth protocol."""
klaude_code/auth/env.py CHANGED
@@ -1,7 +1,7 @@
1
1
  """Environment variable configuration stored in klaude-auth.json."""
2
2
 
3
3
  import json
4
- from typing import Any
4
+ from typing import Any, cast
5
5
 
6
6
  from klaude_code.auth.base import KLAUDE_AUTH_FILE
7
7
 
@@ -11,9 +11,9 @@ def _load_store() -> dict[str, Any]:
11
11
  if not KLAUDE_AUTH_FILE.exists():
12
12
  return {}
13
13
  try:
14
- data: Any = json.loads(KLAUDE_AUTH_FILE.read_text())
14
+ data = json.loads(KLAUDE_AUTH_FILE.read_text())
15
15
  if isinstance(data, dict):
16
- return dict(data)
16
+ return cast(dict[str, Any], data)
17
17
  return {}
18
18
  except (json.JSONDecodeError, ValueError):
19
19
  return {}
@@ -25,26 +25,31 @@ def _save_store(data: dict[str, Any]) -> None:
25
25
  KLAUDE_AUTH_FILE.write_text(json.dumps(data, indent=2))
26
26
 
27
27
 
28
+ def _get_env_section(store: dict[str, Any]) -> dict[str, Any] | None:
29
+ """Extract and validate the 'env' section from store."""
30
+ env_section = store.get("env")
31
+ if isinstance(env_section, dict):
32
+ return cast(dict[str, Any], env_section)
33
+ return None
34
+
35
+
28
36
  def get_auth_env(env_var: str) -> str | None:
29
37
  """Get environment variable value from klaude-auth.json 'env' section.
30
38
 
31
39
  This provides a fallback for API keys when real environment variables are not set.
32
40
  Priority: os.environ > klaude-auth.json env
33
41
  """
34
- store = _load_store()
35
- env_section: Any = store.get("env")
36
- if not isinstance(env_section, dict):
42
+ env_section = _get_env_section(_load_store())
43
+ if env_section is None:
37
44
  return None
38
- value: Any = env_section.get(env_var)
45
+ value = env_section.get(env_var)
39
46
  return str(value) if value is not None else None
40
47
 
41
48
 
42
49
  def set_auth_env(env_var: str, value: str) -> None:
43
50
  """Set environment variable value in klaude-auth.json 'env' section."""
44
51
  store = _load_store()
45
- env_section: Any = store.get("env")
46
- if not isinstance(env_section, dict):
47
- env_section = {}
52
+ env_section = _get_env_section(store) or {}
48
53
  env_section[env_var] = value
49
54
  store["env"] = env_section
50
55
  _save_store(store)
@@ -53,8 +58,8 @@ def set_auth_env(env_var: str, value: str) -> None:
53
58
  def delete_auth_env(env_var: str) -> None:
54
59
  """Delete environment variable from klaude-auth.json 'env' section."""
55
60
  store = _load_store()
56
- env_section: Any = store.get("env")
57
- if not isinstance(env_section, dict):
61
+ env_section = _get_env_section(store)
62
+ if env_section is None:
58
63
  return
59
64
  env_section.pop(env_var, None)
60
65
  if len(env_section) == 0:
@@ -70,8 +75,7 @@ def delete_auth_env(env_var: str) -> None:
70
75
 
71
76
  def list_auth_env() -> dict[str, str]:
72
77
  """List all environment variables in klaude-auth.json 'env' section."""
73
- store = _load_store()
74
- env_section: Any = store.get("env")
75
- if not isinstance(env_section, dict):
78
+ env_section = _get_env_section(_load_store())
79
+ if env_section is None:
76
80
  return {}
77
81
  return {k: str(v) for k, v in env_section.items() if v is not None}
@@ -24,6 +24,11 @@ def _select_provider() -> str | None:
24
24
  value="codex",
25
25
  search_text="codex",
26
26
  ),
27
+ SelectItem(
28
+ title=[("", "Google Antigravity "), ("ansibrightblack", "[OAuth]\n")],
29
+ value="antigravity",
30
+ search_text="antigravity",
31
+ ),
27
32
  ]
28
33
  # Add API key options
29
34
  for key_info in SUPPORTED_API_KEYS:
@@ -71,7 +76,7 @@ def _build_provider_help() -> str:
71
76
  from klaude_code.config.builtin_config import SUPPORTED_API_KEYS
72
77
 
73
78
  # Use first word of name for brevity (e.g., "google" instead of "google gemini")
74
- names = ["codex", "claude"] + [k.name.split()[0].lower() for k in SUPPORTED_API_KEYS]
79
+ names = ["codex", "claude", "antigravity"] + [k.name.split()[0].lower() for k in SUPPORTED_API_KEYS]
75
80
  return f"Provider name ({', '.join(names)})"
76
81
 
77
82
 
@@ -149,6 +154,39 @@ def login_command(
149
154
  except Exception as e:
150
155
  log((f"Login failed: {e}", "red"))
151
156
  raise typer.Exit(1) from None
157
+ case "antigravity":
158
+ from klaude_code.auth.antigravity.oauth import AntigravityOAuth
159
+ from klaude_code.auth.antigravity.token_manager import AntigravityTokenManager
160
+
161
+ token_manager = AntigravityTokenManager()
162
+
163
+ if token_manager.is_logged_in():
164
+ state = token_manager.get_state()
165
+ if state and not state.is_expired():
166
+ log(("You are already logged in to Antigravity.", "green"))
167
+ if state.email:
168
+ log(f" Email: {state.email}")
169
+ log(f" Project ID: {state.project_id}")
170
+ expires_dt = datetime.datetime.fromtimestamp(state.expires_at, tz=datetime.UTC)
171
+ log(f" Expires: {expires_dt.strftime('%Y-%m-%d %H:%M:%S UTC')}")
172
+ if not typer.confirm("Do you want to re-login?"):
173
+ return
174
+
175
+ log("Starting Antigravity OAuth login flow...")
176
+ log("A browser window will open for authentication.")
177
+
178
+ try:
179
+ oauth = AntigravityOAuth(token_manager)
180
+ state = oauth.login()
181
+ log(("Login successful!", "green"))
182
+ if state.email:
183
+ log(f" Email: {state.email}")
184
+ log(f" Project ID: {state.project_id}")
185
+ expires_dt = datetime.datetime.fromtimestamp(state.expires_at, tz=datetime.UTC)
186
+ log(f" Expires: {expires_dt.strftime('%Y-%m-%d %H:%M:%S UTC')}")
187
+ except Exception as e:
188
+ log((f"Login failed: {e}", "red"))
189
+ raise typer.Exit(1) from None
152
190
  case _:
153
191
  from klaude_code.config.builtin_config import SUPPORTED_API_KEYS
154
192
 
@@ -174,7 +212,7 @@ def login_command(
174
212
 
175
213
 
176
214
  def logout_command(
177
- provider: str = typer.Argument("codex", help="Provider to logout (codex|claude)"),
215
+ provider: str = typer.Argument("codex", help="Provider to logout (codex|claude|antigravity)"),
178
216
  ) -> None:
179
217
  """Logout from a provider."""
180
218
  match provider.lower():
@@ -202,8 +240,20 @@ def logout_command(
202
240
  if typer.confirm("Are you sure you want to logout from Claude?"):
203
241
  token_manager.delete()
204
242
  log(("Logged out from Claude.", "green"))
243
+ case "antigravity":
244
+ from klaude_code.auth.antigravity.token_manager import AntigravityTokenManager
245
+
246
+ token_manager = AntigravityTokenManager()
247
+
248
+ if not token_manager.is_logged_in():
249
+ log("You are not logged in to Antigravity.")
250
+ return
251
+
252
+ if typer.confirm("Are you sure you want to logout from Antigravity?"):
253
+ token_manager.delete()
254
+ log(("Logged out from Antigravity.", "green"))
205
255
  case _:
206
- log((f"Error: Unknown provider '{provider}'. Supported: codex, claude", "red"))
256
+ log((f"Error: Unknown provider '{provider}'. Supported: codex, claude, antigravity", "red"))
207
257
  raise typer.Exit(1)
208
258
 
209
259
 
@@ -212,7 +262,7 @@ def register_auth_commands(app: typer.Typer) -> None:
212
262
  auth_app = typer.Typer(help="Login/logout", invoke_without_command=True)
213
263
 
214
264
  @auth_app.callback()
215
- def auth_callback(ctx: typer.Context) -> None:
265
+ def auth_callback(ctx: typer.Context) -> None: # pyright: ignore[reportUnusedFunction]
216
266
  """Authentication commands for managing provider logins."""
217
267
  if ctx.invoked_subcommand is None:
218
268
  typer.echo(ctx.get_help())