Skip to content

Commit cf369d4

Browse files
authored
Merge branch 'router-for-me:main' into main
2 parents e3d8d72 + 3099114 commit cf369d4

File tree

21 files changed

+965
-72
lines changed

21 files changed

+965
-72
lines changed

.dockerignore

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,11 +23,14 @@ config.yaml
2323

2424
# Development/editor
2525
bin/*
26-
.claude/*
2726
.vscode/*
27+
.claude/*
28+
.codex/*
2829
.gemini/*
2930
.serena/*
3031
.agent/*
32+
.agents/*
33+
.opencode/*
3134
.bmad/*
3235
_bmad/*
3336
_bmad-output/*

.gitignore

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,10 +34,14 @@ GEMINI.md
3434

3535
# Tooling metadata
3636
.vscode/*
37+
.codex/*
3738
.claude/*
3839
.gemini/*
3940
.serena/*
4041
.agent/*
42+
.agents/*
43+
.agents/*
44+
.opencode/*
4145
.bmad/*
4246
_bmad/*
4347
_bmad-output/*

config.example.yaml

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ auth-dir: "~/.cli-proxy-api"
3535
api-keys:
3636
- "your-api-key-1"
3737
- "your-api-key-2"
38+
- "your-api-key-3"
3839

3940
# Enable debug logging
4041
debug: false
@@ -181,6 +182,18 @@ ws-auth: false
181182
# upstream-url: "https://ampcode.com"
182183
# # Optional: Override API key for Amp upstream (otherwise uses env or file)
183184
# upstream-api-key: ""
185+
# # Per-client upstream API key mapping
186+
# # Maps client API keys (from top-level api-keys) to different Amp upstream API keys.
187+
# # Useful when different clients need to use different Amp accounts/quotas.
188+
# # If a client key isn't mapped, falls back to upstream-api-key (default behavior).
189+
# upstream-api-keys:
190+
# - upstream-api-key: "amp_key_for_team_a" # Upstream key to use for these clients
191+
# api-keys: # Client keys that use this upstream key
192+
# - "your-api-key-1"
193+
# - "your-api-key-2"
194+
# - upstream-api-key: "amp_key_for_team_b"
195+
# api-keys:
196+
# - "your-api-key-3"
184197
# # Restrict Amp management routes (/api/auth, /api/user, etc.) to localhost only (default: false)
185198
# restrict-management-to-localhost: false
186199
# # Force model mappings to run before checking local API keys (default: false)

internal/api/handlers/management/auth_files.go

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -431,9 +431,46 @@ func (h *Handler) buildAuthFileEntry(auth *coreauth.Auth) gin.H {
431431
log.WithError(err).Warnf("failed to stat auth file %s", path)
432432
}
433433
}
434+
if claims := extractCodexIDTokenClaims(auth); claims != nil {
435+
entry["id_token"] = claims
436+
}
434437
return entry
435438
}
436439

440+
func extractCodexIDTokenClaims(auth *coreauth.Auth) gin.H {
441+
if auth == nil || auth.Metadata == nil {
442+
return nil
443+
}
444+
if !strings.EqualFold(strings.TrimSpace(auth.Provider), "codex") {
445+
return nil
446+
}
447+
idTokenRaw, ok := auth.Metadata["id_token"].(string)
448+
if !ok {
449+
return nil
450+
}
451+
idToken := strings.TrimSpace(idTokenRaw)
452+
if idToken == "" {
453+
return nil
454+
}
455+
claims, err := codex.ParseJWTToken(idToken)
456+
if err != nil || claims == nil {
457+
return nil
458+
}
459+
460+
result := gin.H{}
461+
if v := strings.TrimSpace(claims.CodexAuthInfo.ChatgptAccountID); v != "" {
462+
result["chatgpt_account_id"] = v
463+
}
464+
if v := strings.TrimSpace(claims.CodexAuthInfo.ChatgptPlanType); v != "" {
465+
result["plan_type"] = v
466+
}
467+
468+
if len(result) == 0 {
469+
return nil
470+
}
471+
return result
472+
}
473+
437474
func authEmail(auth *coreauth.Auth) string {
438475
if auth == nil {
439476
return ""

internal/api/handlers/management/config_lists.go

Lines changed: 148 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -940,3 +940,151 @@ func (h *Handler) GetAmpForceModelMappings(c *gin.Context) {
940940
func (h *Handler) PutAmpForceModelMappings(c *gin.Context) {
941941
h.updateBoolField(c, func(v bool) { h.cfg.AmpCode.ForceModelMappings = v })
942942
}
943+
944+
// GetAmpUpstreamAPIKeys returns the ampcode upstream API keys mapping.
945+
func (h *Handler) GetAmpUpstreamAPIKeys(c *gin.Context) {
946+
if h == nil || h.cfg == nil {
947+
c.JSON(200, gin.H{"upstream-api-keys": []config.AmpUpstreamAPIKeyEntry{}})
948+
return
949+
}
950+
c.JSON(200, gin.H{"upstream-api-keys": h.cfg.AmpCode.UpstreamAPIKeys})
951+
}
952+
953+
// PutAmpUpstreamAPIKeys replaces all ampcode upstream API keys mappings.
954+
func (h *Handler) PutAmpUpstreamAPIKeys(c *gin.Context) {
955+
var body struct {
956+
Value []config.AmpUpstreamAPIKeyEntry `json:"value"`
957+
}
958+
if err := c.ShouldBindJSON(&body); err != nil {
959+
c.JSON(400, gin.H{"error": "invalid body"})
960+
return
961+
}
962+
// Normalize entries: trim whitespace, filter empty
963+
normalized := normalizeAmpUpstreamAPIKeyEntries(body.Value)
964+
h.cfg.AmpCode.UpstreamAPIKeys = normalized
965+
h.persist(c)
966+
}
967+
968+
// PatchAmpUpstreamAPIKeys adds or updates upstream API keys entries.
969+
// Matching is done by upstream-api-key value.
970+
func (h *Handler) PatchAmpUpstreamAPIKeys(c *gin.Context) {
971+
var body struct {
972+
Value []config.AmpUpstreamAPIKeyEntry `json:"value"`
973+
}
974+
if err := c.ShouldBindJSON(&body); err != nil {
975+
c.JSON(400, gin.H{"error": "invalid body"})
976+
return
977+
}
978+
979+
existing := make(map[string]int)
980+
for i, entry := range h.cfg.AmpCode.UpstreamAPIKeys {
981+
existing[strings.TrimSpace(entry.UpstreamAPIKey)] = i
982+
}
983+
984+
for _, newEntry := range body.Value {
985+
upstreamKey := strings.TrimSpace(newEntry.UpstreamAPIKey)
986+
if upstreamKey == "" {
987+
continue
988+
}
989+
normalizedEntry := config.AmpUpstreamAPIKeyEntry{
990+
UpstreamAPIKey: upstreamKey,
991+
APIKeys: normalizeAPIKeysList(newEntry.APIKeys),
992+
}
993+
if idx, ok := existing[upstreamKey]; ok {
994+
h.cfg.AmpCode.UpstreamAPIKeys[idx] = normalizedEntry
995+
} else {
996+
h.cfg.AmpCode.UpstreamAPIKeys = append(h.cfg.AmpCode.UpstreamAPIKeys, normalizedEntry)
997+
existing[upstreamKey] = len(h.cfg.AmpCode.UpstreamAPIKeys) - 1
998+
}
999+
}
1000+
h.persist(c)
1001+
}
1002+
1003+
// DeleteAmpUpstreamAPIKeys removes specified upstream API keys entries.
1004+
// Body must be JSON: {"value": ["<upstream-api-key>", ...]}.
1005+
// If "value" is an empty array, clears all entries.
1006+
// If JSON is invalid or "value" is missing/null, returns 400 and does not persist any change.
1007+
func (h *Handler) DeleteAmpUpstreamAPIKeys(c *gin.Context) {
1008+
var body struct {
1009+
Value []string `json:"value"`
1010+
}
1011+
if err := c.ShouldBindJSON(&body); err != nil {
1012+
c.JSON(400, gin.H{"error": "invalid body"})
1013+
return
1014+
}
1015+
1016+
if body.Value == nil {
1017+
c.JSON(400, gin.H{"error": "missing value"})
1018+
return
1019+
}
1020+
1021+
// Empty array means clear all
1022+
if len(body.Value) == 0 {
1023+
h.cfg.AmpCode.UpstreamAPIKeys = nil
1024+
h.persist(c)
1025+
return
1026+
}
1027+
1028+
toRemove := make(map[string]bool)
1029+
for _, key := range body.Value {
1030+
trimmed := strings.TrimSpace(key)
1031+
if trimmed == "" {
1032+
continue
1033+
}
1034+
toRemove[trimmed] = true
1035+
}
1036+
if len(toRemove) == 0 {
1037+
c.JSON(400, gin.H{"error": "empty value"})
1038+
return
1039+
}
1040+
1041+
newEntries := make([]config.AmpUpstreamAPIKeyEntry, 0, len(h.cfg.AmpCode.UpstreamAPIKeys))
1042+
for _, entry := range h.cfg.AmpCode.UpstreamAPIKeys {
1043+
if !toRemove[strings.TrimSpace(entry.UpstreamAPIKey)] {
1044+
newEntries = append(newEntries, entry)
1045+
}
1046+
}
1047+
h.cfg.AmpCode.UpstreamAPIKeys = newEntries
1048+
h.persist(c)
1049+
}
1050+
1051+
// normalizeAmpUpstreamAPIKeyEntries normalizes a list of upstream API key entries.
1052+
func normalizeAmpUpstreamAPIKeyEntries(entries []config.AmpUpstreamAPIKeyEntry) []config.AmpUpstreamAPIKeyEntry {
1053+
if len(entries) == 0 {
1054+
return nil
1055+
}
1056+
out := make([]config.AmpUpstreamAPIKeyEntry, 0, len(entries))
1057+
for _, entry := range entries {
1058+
upstreamKey := strings.TrimSpace(entry.UpstreamAPIKey)
1059+
if upstreamKey == "" {
1060+
continue
1061+
}
1062+
apiKeys := normalizeAPIKeysList(entry.APIKeys)
1063+
out = append(out, config.AmpUpstreamAPIKeyEntry{
1064+
UpstreamAPIKey: upstreamKey,
1065+
APIKeys: apiKeys,
1066+
})
1067+
}
1068+
if len(out) == 0 {
1069+
return nil
1070+
}
1071+
return out
1072+
}
1073+
1074+
// normalizeAPIKeysList trims and filters empty strings from a list of API keys.
1075+
func normalizeAPIKeysList(keys []string) []string {
1076+
if len(keys) == 0 {
1077+
return nil
1078+
}
1079+
out := make([]string, 0, len(keys))
1080+
for _, k := range keys {
1081+
trimmed := strings.TrimSpace(k)
1082+
if trimmed != "" {
1083+
out = append(out, trimmed)
1084+
}
1085+
}
1086+
if len(out) == 0 {
1087+
return nil
1088+
}
1089+
return out
1090+
}

internal/api/handlers/management/handler.go

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,11 @@ func NewHandler(cfg *config.Config, configFilePath string, manager *coreauth.Man
5959
}
6060
}
6161

62+
// NewHandler creates a new management handler instance.
63+
func NewHandlerWithoutConfigFilePath(cfg *config.Config, manager *coreauth.Manager) *Handler {
64+
return NewHandler(cfg, "", manager)
65+
}
66+
6267
// SetConfig updates the in-memory config reference when the server hot-reloads.
6368
func (h *Handler) SetConfig(cfg *config.Config) { h.cfg = cfg }
6469

internal/api/modules/amp/amp.go

Lines changed: 85 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -227,11 +227,20 @@ func (m *AmpModule) OnConfigUpdated(cfg *config.Config) error {
227227
}
228228
}
229229

230-
// Check API key change
230+
// Check API key change (both default and per-client mappings)
231231
apiKeyChanged := m.hasAPIKeyChanged(oldSettings, &newSettings)
232-
if apiKeyChanged {
232+
upstreamAPIKeysChanged := m.hasUpstreamAPIKeysChanged(oldSettings, &newSettings)
233+
if apiKeyChanged || upstreamAPIKeysChanged {
233234
if m.secretSource != nil {
234-
if ms, ok := m.secretSource.(*MultiSourceSecret); ok {
235+
if ms, ok := m.secretSource.(*MappedSecretSource); ok {
236+
if apiKeyChanged {
237+
ms.UpdateDefaultExplicitKey(newSettings.UpstreamAPIKey)
238+
ms.InvalidateCache()
239+
}
240+
if upstreamAPIKeysChanged {
241+
ms.UpdateMappings(newSettings.UpstreamAPIKeys)
242+
}
243+
} else if ms, ok := m.secretSource.(*MultiSourceSecret); ok {
235244
ms.UpdateExplicitKey(newSettings.UpstreamAPIKey)
236245
ms.InvalidateCache()
237246
}
@@ -251,10 +260,22 @@ func (m *AmpModule) OnConfigUpdated(cfg *config.Config) error {
251260

252261
func (m *AmpModule) enableUpstreamProxy(upstreamURL string, settings *config.AmpCode) error {
253262
if m.secretSource == nil {
254-
m.secretSource = NewMultiSourceSecret(settings.UpstreamAPIKey, 0 /* default 5min */)
263+
// Create MultiSourceSecret as the default source, then wrap with MappedSecretSource
264+
defaultSource := NewMultiSourceSecret(settings.UpstreamAPIKey, 0 /* default 5min */)
265+
mappedSource := NewMappedSecretSource(defaultSource)
266+
mappedSource.UpdateMappings(settings.UpstreamAPIKeys)
267+
m.secretSource = mappedSource
268+
} else if ms, ok := m.secretSource.(*MappedSecretSource); ok {
269+
ms.UpdateDefaultExplicitKey(settings.UpstreamAPIKey)
270+
ms.InvalidateCache()
271+
ms.UpdateMappings(settings.UpstreamAPIKeys)
255272
} else if ms, ok := m.secretSource.(*MultiSourceSecret); ok {
273+
// Legacy path: wrap existing MultiSourceSecret with MappedSecretSource
256274
ms.UpdateExplicitKey(settings.UpstreamAPIKey)
257275
ms.InvalidateCache()
276+
mappedSource := NewMappedSecretSource(ms)
277+
mappedSource.UpdateMappings(settings.UpstreamAPIKeys)
278+
m.secretSource = mappedSource
258279
}
259280

260281
proxy, err := createReverseProxy(upstreamURL, m.secretSource)
@@ -313,6 +334,66 @@ func (m *AmpModule) hasAPIKeyChanged(old *config.AmpCode, new *config.AmpCode) b
313334
return oldKey != newKey
314335
}
315336

337+
// hasUpstreamAPIKeysChanged compares old and new per-client upstream API key mappings.
338+
func (m *AmpModule) hasUpstreamAPIKeysChanged(old *config.AmpCode, new *config.AmpCode) bool {
339+
if old == nil {
340+
return len(new.UpstreamAPIKeys) > 0
341+
}
342+
343+
if len(old.UpstreamAPIKeys) != len(new.UpstreamAPIKeys) {
344+
return true
345+
}
346+
347+
// Build map for comparison: upstreamKey -> set of clientKeys
348+
type entryInfo struct {
349+
upstreamKey string
350+
clientKeys map[string]struct{}
351+
}
352+
oldEntries := make([]entryInfo, len(old.UpstreamAPIKeys))
353+
for i, entry := range old.UpstreamAPIKeys {
354+
clientKeys := make(map[string]struct{}, len(entry.APIKeys))
355+
for _, k := range entry.APIKeys {
356+
trimmed := strings.TrimSpace(k)
357+
if trimmed == "" {
358+
continue
359+
}
360+
clientKeys[trimmed] = struct{}{}
361+
}
362+
oldEntries[i] = entryInfo{
363+
upstreamKey: strings.TrimSpace(entry.UpstreamAPIKey),
364+
clientKeys: clientKeys,
365+
}
366+
}
367+
368+
for i, newEntry := range new.UpstreamAPIKeys {
369+
if i >= len(oldEntries) {
370+
return true
371+
}
372+
oldE := oldEntries[i]
373+
if strings.TrimSpace(newEntry.UpstreamAPIKey) != oldE.upstreamKey {
374+
return true
375+
}
376+
newKeys := make(map[string]struct{}, len(newEntry.APIKeys))
377+
for _, k := range newEntry.APIKeys {
378+
trimmed := strings.TrimSpace(k)
379+
if trimmed == "" {
380+
continue
381+
}
382+
newKeys[trimmed] = struct{}{}
383+
}
384+
if len(newKeys) != len(oldE.clientKeys) {
385+
return true
386+
}
387+
for k := range newKeys {
388+
if _, ok := oldE.clientKeys[k]; !ok {
389+
return true
390+
}
391+
}
392+
}
393+
394+
return false
395+
}
396+
316397
// GetModelMapper returns the model mapper instance (for testing/debugging).
317398
func (m *AmpModule) GetModelMapper() *DefaultModelMapper {
318399
return m.modelMapper

0 commit comments

Comments
 (0)