Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/client/cli/commands/up.py: 0%

1import atexit 

2import contextlib 

3import json 

4import os 

5import signal 

6import sys 

7import threading 

8from collections.abc import Iterator 

9from dataclasses import dataclass 

10from pathlib import Path 

11from types import FrameType 

12from typing import IO, Final 

13 

14import click 

15from pydantic import JsonValue, TypeAdapter, ValidationError 

16 

17from litellm.litellm_core_utils.cli_keyring import SecretVault 

18from litellm.litellm_core_utils.cli_token_utils import is_cli_token_fresh 

19from litellm.litellm_core_utils.private_json import ensure_private_dir 

20 

21from .agents import AgentRunError, resolve_api_key, verify_proxy_key 

22from .auth import CliContextObj, context_secret_vault, get_stored_api_key, load_token, login 

23from .claude_settings import ( 

24 BACKUP_PATH, 

25 CLAUDE_SETTINGS_PATH, 

26 ClaudeSettingsError, 

27 StaticToken, 

28 install_statusline_script, 

29 load_json_or_empty, 

30 merge_claude_settings, 

31 write_claude_settings, 

32) 

33 

34 

35class UpError(ClaudeSettingsError): 

36 """Raised for any user-actionable failure while starting/stopping interception.""" 

37 

38 

39@dataclass(frozen=True, slots=True) 

40class BackupRecord: 

41 """Snapshot of ~/.claude/settings.json taken right before `lite up` patches it.""" 

42 

43 existed: bool 

44 content: dict[str, JsonValue] | None 

45 

46 

47_BACKUP_RECORD_ADAPTER: Final = TypeAdapter(BackupRecord) 

48 

49 

50@contextlib.contextmanager 

51def secure_create(path: Path) -> Iterator[IO[str]]: 

52 """Open path for writing with mode 0600 fixed up before any content is written. 

53 

54 A plain `open(path, "w")` creates a *new* file at the umask-derived default (commonly 0644) 

55 and leaves it world- or group-readable until a later `chmod` call catches up -- a real window 

56 in which a file holding a credential is readable by another local account. Passing the mode to 

57 `os.open` closes that window for a brand-new file, but `O_CREAT`'s mode argument is only 

58 applied on creation: if the file already exists its old, broader permissions carry over 

59 untouched. `os.fchmod` right after opening -- before a single byte of the new content is 

60 written -- covers both cases. 

61 """ 

62 fd: Final = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) 

63 os.fchmod(fd, 0o600) 

64 f: Final[IO[str]] = os.fdopen(fd, "w") 

65 try: 

66 yield f 

67 finally: 

68 f.close() 

69 

70 

71def write_backup(record: BackupRecord, backup_path: Path | None = None) -> None: 

72 path: Final = backup_path if backup_path is not None else BACKUP_PATH 

73 ensure_private_dir(path.parent) 

74 with secure_create(path) as f: 

75 json.dump({"existed": record.existed, "content": record.content}, f, indent=2) 

76 

77 

78def read_backup(backup_path: Path | None = None) -> BackupRecord | None: 

79 path: Final = backup_path if backup_path is not None else BACKUP_PATH 

80 if not path.exists(): 

81 return None 

82 with open(path, "r") as f: 

83 content: Final = f.read() 

84 try: 

85 return _BACKUP_RECORD_ADAPTER.validate_json(content) 

86 except ValidationError: 

87 raise UpError(f"{path} contains invalid or unexpected JSON; cannot restore from it safely.") 

88 

89 

90def restore_claude_settings(settings_path: Path | None = None, backup_path: Path | None = None) -> BackupRecord | None: 

91 """Restore settings_path from the backup at backup_path, then delete the backup. 

92 

93 Returns the restored record, or None if there was nothing to restore. 

94 """ 

95 resolved_settings_path: Final = settings_path if settings_path is not None else CLAUDE_SETTINGS_PATH 

96 resolved_backup_path: Final = backup_path if backup_path is not None else BACKUP_PATH 

97 record: Final = read_backup(resolved_backup_path) 

98 if record is None: 

99 return None 

100 if record.existed and record.content is not None: 

101 resolved_settings_path.parent.mkdir(parents=True, exist_ok=True) 

102 write_claude_settings(resolved_settings_path, record.content) 

103 elif resolved_settings_path.exists(): 

104 resolved_settings_path.unlink() 

105 resolved_backup_path.unlink() 

106 return record 

107 

108 

109def _usable_login(api_key: str | None, vault: SecretVault) -> bool: 

110 if api_key is None: 

111 return False 

112 token_data: Final = load_token(vault=vault) 

113 return token_data is not None and is_cli_token_fresh(token_data) 

114 

115 

116def _key_resolved_on_the_way_in(ctx_obj: CliContextObj, base_url: str, vault: SecretVault) -> str | None: 

117 if ctx_obj.get("api_key_from_token_file"): 

118 return ctx_obj.get("api_key") 

119 return get_stored_api_key(expected_base_url=base_url, vault=vault) 

120 

121 

122def _stored_login_is_pkce(vault: SecretVault) -> bool: 

123 token_data: Final = load_token(vault=vault) 

124 return token_data is not None and token_data.get("refresh_token") is not None 

125 

126 

127def ensure_fresh_login(ctx: click.Context) -> None: 

128 ctx_obj: Final[CliContextObj] = ctx.obj 

129 base_url: Final = ctx_obj["base_url"].rstrip("/") 

130 vault: Final = context_secret_vault(ctx) 

131 if _usable_login(_key_resolved_on_the_way_in(ctx_obj, base_url, vault), vault): 

132 return 

133 

134 pkce: Final = _stored_login_is_pkce(vault) 

135 login_command: Final = "lite login --pkce" if pkce else "lite login" 

136 if not sys.stdin.isatty(): 

137 raise UpError(f"No fresh LiteLLM login found for this proxy. Run `{login_command}` first.") 

138 

139 click.echo("No fresh LiteLLM login found for this proxy; starting login...") 

140 ctx.invoke(login, config_claude=False, pkce=pkce) 

141 if not _usable_login(get_stored_api_key(expected_base_url=base_url, vault=vault), vault): 

142 raise UpError("Login did not produce a usable token.") 

143 

144 

145def _restore_and_report() -> None: 

146 record: Final = restore_claude_settings() 

147 if record is None: 

148 click.echo("Nothing to restore.") 

149 return 

150 if record.existed: 

151 click.echo(f"Restored {CLAUDE_SETTINGS_PATH} to its original contents.") 

152 else: 

153 click.echo(f"Removed {CLAUDE_SETTINGS_PATH} (it did not exist before `lite up`).") 

154 

155 

156@click.command(name="up") 

157@click.pass_context 

158def up(ctx: click.Context) -> None: 

159 """Route every Claude Code session through your LiteLLM proxy until stopped. 

160 

161 Patches ~/.claude/settings.json so Claude Code picks up the proxy on its own 

162 next startup, from any terminal -- no need to launch it through `lite`. The 

163 key written is the one this command resolved (your fresh `lite login`, or an 

164 explicit --api-key), copied in as a static token for as long as `up` runs. 

165 Press Ctrl-C to stop and restore your original settings. Assumes the proxy 

166 is already running (this does not start one for you). Cursor is not 

167 supported: it has no equivalent file-based config to patch. 

168 """ 

169 ctx_obj: Final[CliContextObj] = ctx.obj 

170 base_url: Final = ctx_obj["base_url"] 

171 

172 try: 

173 ensure_fresh_login(ctx) 

174 api_key: Final = resolve_api_key(ctx) 

175 verify_proxy_key(base_url, api_key) 

176 

177 if BACKUP_PATH.exists(): 

178 raise UpError( 

179 f"{BACKUP_PATH} already exists -- `lite up` looks like it's already " 

180 "running (or crashed without cleanup). Run `lite down` first." 

181 ) 

182 

183 status_line: Final = install_statusline_script() 

184 original_existed: Final = CLAUDE_SETTINGS_PATH.exists() 

185 original_settings: Final = load_json_or_empty(CLAUDE_SETTINGS_PATH) 

186 write_backup( 

187 BackupRecord( 

188 existed=original_existed, 

189 content=original_settings if original_existed else None, 

190 ) 

191 ) 

192 

193 CLAUDE_SETTINGS_PATH.parent.mkdir(exist_ok=True) 

194 merged: Final = merge_claude_settings( 

195 original_settings, base_url, StaticToken(api_key), status_line=status_line 

196 ) 

197 write_claude_settings(CLAUDE_SETTINGS_PATH, merged) 

198 except (AgentRunError, ClaudeSettingsError) as e: 

199 raise click.ClickException(str(e)) 

200 

201 click.echo(f"litellm: routing Claude Code through proxy at {base_url.rstrip('/')}") 

202 click.echo("Press Ctrl-C to stop and restore your original settings.") 

203 

204 stop_event: Final = threading.Event() 

205 restored: Final = threading.Lock() 

206 

207 def _handle_signal(_signum: int, _frame: FrameType | None) -> None: 

208 stop_event.set() 

209 

210 def _restore_once() -> None: 

211 if not restored.acquire(blocking=False): 

212 return 

213 try: 

214 _restore_and_report() 

215 except ClaudeSettingsError as e: 

216 # Runs from atexit/a signal handler, outside Click's own exception 

217 # handling -- raising here would only produce an unhandled-exception 

218 # warning on stderr, not a clean message. 

219 click.echo(str(e), err=True) 

220 

221 signal.signal(signal.SIGINT, _handle_signal) 

222 signal.signal(signal.SIGTERM, _handle_signal) 

223 atexit.register(_restore_once) 

224 

225 stop_event.wait() 

226 _restore_once() 

227 

228 

229@click.command(name="down") 

230def down() -> None: 

231 """Restore ~/.claude/settings.json if a `lite up` session left it patched. 

232 

233 Use this after a `lite up` process was killed uncleanly (e.g. `kill -9`) 

234 instead of stopped with Ctrl-C. 

235 """ 

236 try: 

237 _restore_and_report() 

238 except ClaudeSettingsError as e: 

239 raise click.ClickException(str(e)) 

240 

241 

242__all__ = [ 

243 "BACKUP_PATH", 

244 "CLAUDE_SETTINGS_PATH", 

245 "BackupRecord", 

246 "ClaudeSettingsError", 

247 "UpError", 

248 "down", 

249 "load_json_or_empty", 

250 "merge_claude_settings", 

251 "read_backup", 

252 "restore_claude_settings", 

253 "up", 

254 "write_backup", 

255]