History

fix(meta): exclude cleared expired timestamp on persist, add runtime alias resolution, and fix test race

Kenny committed 2 months ago

Commit 47cc31a1 parent cee799f8 files · +251 −3Browse files

Unified · Split

8 files
  1. internal/auth/meta/meta.go +22 −0
  2. internal/auth/meta/meta_auth_test.go +127 −0
  3. internal/runtime/executor/meta_executor_test.go +2 −2
  4. sdk/cliproxy/auth/api_key_model_alias_test.go +66 −0
  5. sdk/cliproxy/auth/api_key_model_capabilities.go +4 −0
  6. sdk/cliproxy/auth/conductor_execution.go +1 −1
  7. sdk/cliproxy/auth/conductor_models.go +25 −0
  8. sdk/cliproxy/auth/conductor_request_scoped_errors.go +4 −0
internal/auth/meta/meta.go+22−0
@@ -111,6 +111,22 @@ func (ts *MetaTokenStorage) SetMetadata(meta map[string]any) { · 110 unmodified lines
111111 ts.Metadata = meta
112112 }
113113
114+// metaCredentialFields identifies credential and lifecycle fields managed by MetaTokenStorage.
115+// These fields must never be restored from disk or metadata when omitted from storage.
116+var metaCredentialFields = map[string]struct{}{
117+ "type": {},
118+ "auth_kind": {},
119+ "access_token": {},
120+ "token_type": {},
121+ "dca_token": {},
122+ "api_key": {},
123+ "expires_in": {},
124+ "expired": {},
125+ "dca_expired": {},
126+ "dca_expires_at": {},
127+ "last_refresh": {},
128+}
129+
114130 // SaveTokenToFile writes Meta credentials to a JSON auth file.
115131 func (ts *MetaTokenStorage) SaveTokenToFile(authFilePath string) error {
116132 ts.Type = "meta"
@@ -161,6 +177,9 @@ func (ts *MetaTokenStorage) SaveTokenToFile(authFilePath string) error { · 44 unmodified lines
161177 var existing map[string]any
162178 if errJSON := json.Unmarshal(raw, &existing); errJSON == nil {
163179 for k, v := range existing {
180+ if _, isCred := metaCredentialFields[k]; isCred {
181+ continue
182+ }
164183 if _, exists := data[k]; !exists {
165184 data[k] = v
166185 }
@@ -168,6 +187,9 @@ func (ts *MetaTokenStorage) SaveTokenToFile(authFilePath string) error { · 1 unmodified line
168187 }
169188 }
170189 for k, v := range ts.Metadata {
190+ if _, isCred := metaCredentialFields[k]; isCred {
191+ continue
192+ }
171193 if _, exists := data[k]; !exists {
172194 data[k] = v
173195 }
internal/auth/meta/meta_auth_test.go+127−0
@@ -190,3 +190,130 @@ func TestReadLocalMuseCLIAuth(t *testing.T) { · 189 unmodified lines
190190 t.Errorf("expected email engineer@meta.com, got %s", cred.Email)
191191 }
192192 }
193+
194+func TestSaveTokenToFile_PreservesCustomMetadata(t *testing.T) {
195+ tempDir := t.TempDir()
196+ authFilePath := filepath.Join(tempDir, "meta-test.json")
197+
198+ storage := &MetaTokenStorage{
199+ Type: "meta",
200+ AuthKind: "oauth",
201+ AccessToken: "test-access-token",
202+ APIKey: "test-api-key",
203+ DCAToken: "dca:token",
204+ Email: "user@example.com",
205+ }
206+ storage.SetMetadata(map[string]any{
207+ "priority": float64(10),
208+ "disable_cooling": true,
209+ "proxy_url": "http://proxy:8080",
210+ })
211+
212+ if errSave := storage.SaveTokenToFile(authFilePath); errSave != nil {
213+ t.Fatalf("SaveTokenToFile() error = %v", errSave)
214+ }
215+
216+ savedRaw, errRead := os.ReadFile(authFilePath)
217+ if errRead != nil {
218+ t.Fatalf("os.ReadFile error = %v", errRead)
219+ }
220+
221+ var saved map[string]any
222+ if errUnmarshal := json.Unmarshal(savedRaw, &saved); errUnmarshal != nil {
223+ t.Fatalf("json.Unmarshal error = %v", errUnmarshal)
224+ }
225+
226+ if saved["api_key"] != "test-api-key" {
227+ t.Errorf("api_key = %v, want test-api-key", saved["api_key"])
228+ }
229+ if saved["priority"] != float64(10) {
230+ t.Errorf("priority = %v, want 10", saved["priority"])
231+ }
232+ if saved["disable_cooling"] != true {
233+ t.Errorf("disable_cooling = %v, want true", saved["disable_cooling"])
234+ }
235+ if saved["proxy_url"] != "http://proxy:8080" {
236+ t.Errorf("proxy_url = %v, want http://proxy:8080", saved["proxy_url"])
237+ }
238+}
239+
240+func TestSaveTokenToFile_DoesNotRestoreClearedCredentialFields(t *testing.T) {
241+ tempDir := t.TempDir()
242+ authFilePath := filepath.Join(tempDir, "meta-test.json")
243+
244+ // Pre-populate an auth file that had expired DCA credentials and custom configuration
245+ initial := map[string]any{
246+ "type": "meta",
247+ "auth_kind": "oauth",
248+ "access_token": "old-dca-token",
249+ "dca_token": "old-dca-token",
250+ "expired": "2026-01-01T00:00:00Z",
251+ "dca_expired": "2026-01-01T00:00:00Z",
252+ "dca_expires_at": float64(1767225600),
253+ "priority": float64(42),
254+ "models": []any{"muse-latest"},
255+ "disable-cooling": true,
256+ }
257+ raw, err := json.Marshal(initial)
258+ if err != nil {
259+ t.Fatalf("json.Marshal error = %v", err)
260+ }
261+ if err := os.WriteFile(authFilePath, raw, 0600); err != nil {
262+ t.Fatalf("WriteFile error = %v", err)
263+ }
264+
265+ // Token storage after minting an API key: Expired is explicitly empty, new APIKey set
266+ storage := &MetaTokenStorage{
267+ Type: "meta",
268+ AuthKind: "oauth",
269+ AccessToken: "new-minted-key",
270+ APIKey: "new-minted-key",
271+ DCAToken: "new-dca-token",
272+ Expired: "", // Deliberately cleared on refresh
273+ Email: "user@example.com",
274+ }
275+
276+ if errSave := storage.SaveTokenToFile(authFilePath); errSave != nil {
277+ t.Fatalf("SaveTokenToFile error = %v", errSave)
278+ }
279+
280+ savedRaw, errRead := os.ReadFile(authFilePath)
281+ if errRead != nil {
282+ t.Fatalf("os.ReadFile error = %v", errRead)
283+ }
284+
285+ var saved map[string]any
286+ if errUnmarshal := json.Unmarshal(savedRaw, &saved); errUnmarshal != nil {
287+ t.Fatalf("json.Unmarshal error = %v", errUnmarshal)
288+ }
289+
290+ // Verify cleared credential fields were NOT restored from disk
291+ if val, exists := saved["expired"]; exists && val != "" {
292+ t.Errorf("SaveTokenToFile restored expired timestamp: %v", val)
293+ }
294+ if val, exists := saved["dca_expired"]; exists && val != "" {
295+ t.Errorf("SaveTokenToFile restored dca_expired timestamp: %v", val)
296+ }
297+ if _, exists := saved["dca_expires_at"]; exists {
298+ t.Errorf("SaveTokenToFile restored dca_expires_at timestamp")
299+ }
300+
301+ // Verify new credentials were written
302+ if saved["api_key"] != "new-minted-key" {
303+ t.Errorf("api_key = %v, want new-minted-key", saved["api_key"])
304+ }
305+ if saved["access_token"] != "new-minted-key" {
306+ t.Errorf("access_token = %v, want new-minted-key", saved["access_token"])
307+ }
308+
309+ // Verify non-credential configuration from disk was preserved
310+ if saved["priority"] != float64(42) {
311+ t.Errorf("priority = %v, want 42", saved["priority"])
312+ }
313+ if saved["disable-cooling"] != true {
314+ t.Errorf("disable-cooling = %v, want true", saved["disable-cooling"])
315+ }
316+ if models, ok := saved["models"].([]any); !ok || len(models) != 1 || models[0] != "muse-latest" {
317+ t.Errorf("models = %v, want [muse-latest]", saved["models"])
318+ }
319+}
internal/runtime/executor/meta_executor_test.go+2−2
@@ -306,7 +306,7 @@ func TestMetaExecutor_Refresh_SingleflightAndMultiAccount(t *testing.T) { · 305 unmodified lines
306306 wg.Add(1)
307307 go func() {
308308 defer wg.Done()
309− res, err := exec.Refresh(context.Background(), auth1)
309+ res, err := exec.Refresh(context.Background(), auth1.Clone())
310310 if err != nil {
311311 t.Errorf("Refresh acct1 error: %v", err)
312312 }
@@ -326,7 +326,7 @@ func TestMetaExecutor_Refresh_SingleflightAndMultiAccount(t *testing.T) { · 13 unmodified lines
326326 Provider: "meta",
327327 Metadata: map[string]any{"dca_token": "dca:acct2"},
328328 }
329− res2, err2 := exec.Refresh(context.Background(), auth2)
329+ res2, err2 := exec.Refresh(context.Background(), auth2.Clone())
330330 if err2 != nil {
331331 t.Fatalf("Refresh acct2 error: %v", err2)
332332 }
sdk/cliproxy/auth/api_key_model_alias_test.go+66−0
@@ -291,3 +291,69 @@ func TestResolveAPIKeyModelAliasWithResult_ForceMappingUsesConfigAliasNotRequest · 290 unmodified lines
291291 t.Fatalf("OriginalAlias = %q want claude-sonnet-4-5", result.OriginalAlias)
292292 }
293293 }
294+
295+func TestLookupAPIKeyUpstreamModel_MetaKey(t *testing.T) {
296+ cfg := &internalconfig.Config{
297+ MetaKey: []internalconfig.MetaKey{
298+ {
299+ APIKey: "meta-key",
300+ BaseURL: "https://api.meta.ai/v1",
301+ Models: []internalconfig.CodexModel{
302+ {Name: "muse-spark-1.3", Alias: "muse-latest"},
303+ },
304+ },
305+ },
306+ }
307+
308+ mgr := NewManager(nil, nil, nil)
309+ mgr.SetConfig(cfg)
310+
311+ ctx := context.Background()
312+ auth := &Auth{
313+ ID: "meta-auth-1",
314+ Provider: "meta",
315+ Attributes: map[string]string{
316+ "api_key": "meta-key",
317+ "base_url": "https://api.meta.ai/v1",
318+ "auth_kind": "apikey",
319+ },
320+ }
321+ if _, err := mgr.Register(ctx, auth); err != nil {
322+ t.Fatalf("register auth: %v", err)
323+ }
324+
325+ // 1. Fast path: lookup per-auth mapping table compiled during register.
326+ resolved := mgr.lookupAPIKeyUpstreamModel("meta-auth-1", "muse-latest")
327+ if resolved != "muse-spark-1.3" {
328+ t.Fatalf("lookupAPIKeyUpstreamModel() = %q, want muse-spark-1.3", resolved)
329+ }
330+
331+ // 2. Slow path: directly call mgr.applyAPIKeyModelAliasWithRouting with an empty alias table to exercise config resolution fallback.
332+ slowRouting := &apiKeyModelRoutingSnapshot{
333+ config: cfg,
334+ aliases: make(apiKeyModelAliasTable),
335+ }
336+ slowResolved := mgr.applyAPIKeyModelAliasWithRouting(slowRouting, auth, "muse-latest")
337+ if slowResolved != "muse-spark-1.3" {
338+ t.Fatalf("applyAPIKeyModelAliasWithRouting(slow) = %q, want muse-spark-1.3", slowResolved)
339+ }
340+
341+ // 3. Model alias result with force mapping / alias metadata
342+ aliasResult := mgr.resolveAPIKeyModelAliasWithResult(auth, "muse-latest")
343+ if aliasResult.UpstreamModel != "muse-spark-1.3" {
344+ t.Fatalf("resolveAPIKeyModelAliasWithResult() upstream = %q, want muse-spark-1.3", aliasResult.UpstreamModel)
345+ }
346+
347+ // 4. Configured alias entries helper
348+ entries := configuredModelAliasEntries(cfg, auth)
349+ found := false
350+ for _, e := range entries {
351+ if e.GetAlias() == "muse-latest" && e.GetName() == "muse-spark-1.3" {
352+ found = true
353+ break
354+ }
355+ }
356+ if !found {
357+ t.Fatalf("configuredModelAliasEntries did not contain muse-latest -> muse-spark-1.3: %+v", entries)
358+ }
359+}
sdk/cliproxy/auth/api_key_model_capabilities.go+4−0
@@ -202,6 +202,10 @@ func compileAPIKeyModelCapabilitiesForAuth(cfg *internalconfig.Config, auth *Aut · 201 unmodified lines
202202 if entry := resolveVertexAPIKeyConfig(cfg, auth); entry != nil {
203203 compileConfiguredModelCapabilities(out, entry.Models, "gemini")
204204 }
205+ case "meta":
206+ if entry := resolveMetaAPIKeyConfig(cfg, auth); entry != nil {
207+ compileConfiguredModelCapabilities(out, entry.Models, "meta")
208+ }
205209 default:
206210 providerKey, compatName := "", ""
207211 if auth.Attributes != nil {
sdk/cliproxy/auth/conductor_execution.go+1−1
@@ -378,7 +378,7 @@ func requestToFormat(provider string, executor ProviderExecutor, req cliproxyexe · 377 unmodified lines
378378 return sdktranslator.FormatClaude
379379 case "gemini", "vertex", "aistudio":
380380 return sdktranslator.FormatGemini
381− case "kimi":
381+ case "kimi", "meta":
382382 return sdktranslator.FormatOpenAI
383383 case "antigravity":
384384 return sdktranslator.FormatAntigravity
sdk/cliproxy/auth/conductor_models.go+25−0
@@ -465,6 +465,10 @@ func configuredModelAliasEntries(cfg *internalconfig.Config, auth *Auth) []model · 464 unmodified lines
465465 if entry := resolveVertexAPIKeyConfig(cfg, auth); entry != nil {
466466 models = asModelAliasEntries(entry.Models)
467467 }
468+ case "meta":
469+ if entry := resolveMetaAPIKeyConfig(cfg, auth); entry != nil {
470+ models = asModelAliasEntries(entry.Models)
471+ }
468472 default:
469473 providerKey := ""
470474 compatName := ""
@@ -621,6 +625,10 @@ func (m *Manager) rebuildAPIKeyModelAliasLocked(cfg *internalconfig.Config) { · 150 unmodified lines
621625 if entry := resolveVertexAPIKeyConfig(cfg, auth); entry != nil {
622626 compileAPIKeyModelAliasForModels(byAlias, entry.Models)
623627 }
628+ case "meta":
629+ if entry := resolveMetaAPIKeyConfig(cfg, auth); entry != nil {
630+ compileAPIKeyModelAliasForModels(byAlias, entry.Models)
631+ }
624632 default:
625633 // OpenAI-compat uses config selection from auth.Attributes.
626634 providerKey := ""
@@ -741,6 +749,8 @@ func (m *Manager) applyAPIKeyModelAliasWithRouting(routing *apiKeyModelRoutingSn · 114 unmodified lines
741749 upstreamModel = resolveUpstreamModelForXAIAPIKey(cfg, auth, requestedModel)
742750 case "vertex":
743751 upstreamModel = resolveUpstreamModelForVertexAPIKey(cfg, auth, requestedModel)
752+ case "meta":
753+ upstreamModel = resolveUpstreamModelForMetaAPIKey(cfg, auth, requestedModel)
744754 default:
745755 upstreamModel = resolveUpstreamModelForOpenAICompatAPIKey(cfg, auth, requestedModel)
746756 }
@@ -848,6 +858,13 @@ func resolveVertexAPIKeyConfig(cfg *internalconfig.Config, auth *Auth) *internal · 101 unmodified lines
848858 return resolveAPIKeyConfig(cfg.VertexCompatAPIKey, auth)
849859 }
850860
861+func resolveMetaAPIKeyConfig(cfg *internalconfig.Config, auth *Auth) *internalconfig.MetaKey {
862+ if cfg == nil {
863+ return nil
864+ }
865+ return resolveAPIKeyConfig(cfg.MetaKey, auth)
866+}
867+
851868 func resolveUpstreamModelForGeminiAPIKey(cfg *internalconfig.Config, auth *Auth, requestedModel string) string {
852869 entry := resolveGeminiAPIKeyConfig(cfg, auth)
853870 if entry == nil {
@@ -896,6 +913,14 @@ func resolveUpstreamModelForVertexAPIKey(cfg *internalconfig.Config, auth *Auth, · 42 unmodified lines
896913 return resolveModelAliasFromConfigModels(requestedModel, asModelAliasEntries(entry.Models))
897914 }
898915
916+func resolveUpstreamModelForMetaAPIKey(cfg *internalconfig.Config, auth *Auth, requestedModel string) string {
917+ entry := resolveMetaAPIKeyConfig(cfg, auth)
918+ if entry == nil {
919+ return ""
920+ }
921+ return resolveModelAliasFromConfigModels(requestedModel, asModelAliasEntries(entry.Models))
922+}
923+
899924 func resolveUpstreamModelForOpenAICompatAPIKey(cfg *internalconfig.Config, auth *Auth, requestedModel string) string {
900925 providerKey := ""
901926 compatName := ""
sdk/cliproxy/auth/conductor_request_scoped_errors.go+4−0
@@ -143,6 +143,10 @@ func extractRequestScopedErrorRules(auth *Auth, cfg *internalconfig.Config) []in · 142 unmodified lines
143143 if index >= 0 && index < len(cfg.XAIKey) {
144144 return cfg.XAIKey[index].RequestScopedErrors
145145 }
146+ case "meta":
147+ if index >= 0 && index < len(cfg.MetaKey) {
148+ return cfg.MetaKey[index].RequestScopedErrors
149+ }
146150 case "gemini":
147151 if index >= 0 && index < len(cfg.GeminiKey) {
148152 return cfg.GeminiKey[index].RequestScopedErrors

Keyboard shortcuts

General

Show keyboard shortcuts
?
Focus the current search
/

Issues and pull requests

Open filters
f

Issues, pull requests and Inbox

Select next item
j↓
Select previous item
k↑
Open selected item
Enter
Dismiss filters or clear selection
Esc

Repository code

Go to file
t