package server import ( "context" "encoding/json" "errors" "fmt" "iter" "sync" "testing" "time" "github.com/Autumn-27/artex/agent" "github.com/Autumn-27/artex/db" "github.com/Autumn-27/norma/harness" "github.com/Autumn-27/norma/llm" ) type scriptedLLMProvider struct { events []llm.StreamEvent err error calls int } func (p *scriptedLLMProvider) Stream(context.Context, llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] { return func(yield func(llm.StreamEvent, error) bool) { p.calls++ for _, event := range p.events { if !yield(event, nil) { return } } if p.err != nil { yield(llm.StreamEvent{}, p.err) } } } func collectTaskLLMStream(seq iter.Seq2[llm.StreamEvent, error]) ([]llm.StreamEvent, error) { var events []llm.StreamEvent var streamErr error for event, err := range seq { if err != nil { streamErr = err break } events = append(events, event) } return events, streamErr } func TestIsQuotaExhaustedError(t *testing.T) { t.Parallel() positive := []string{ `openai: code=insufficient_quota status=429`, `provider error: quota_exceeded`, `grpc code=RESOURCE_EXHAUSTED`, `You exceeded your current quota, please check your plan and billing details.`, `billing_hard_limit_reached`, `billing_not_active`, `HTTP status 402 Payment Required`, `账户余额不足,请充值`, `credit balance is too low`, } for _, message := range positive { if !isQuotaExhaustedError(errors.New(message)) { t.Errorf("expected quota classification for %q", message) } } negative := []string{ `HTTP status 429: rate limit exceeded`, `RESOURCE_EXHAUSTED: rate limit exceeded`, `HTTP status 429: quota exceeded for quota metric GenerateRequestsPerMinutePerProjectPerBaseModel`, `RESOURCE_EXHAUSTED: TPM quota exceeded`, `HTTP status 401: invalid api key`, `HTTP status 401: insufficient_quota`, `HTTP status 403: forbidden`, `HTTP status 403: billing_hard_limit_reached`, `HTTP status 500: internal server error`, `HTTP status 503: quota_exceeded`, `dial tcp: network is unreachable`, `context length exceeded`, `invalid request parameter`, } for _, message := range negative { if isQuotaExhaustedError(errors.New(message)) { t.Errorf("unexpected quota classification for %q", message) } } if isQuotaExhaustedError(nil) { t.Error("nil error must not be classified as quota exhaustion") } } func TestTaskDTOUsesEmptyArrays(t *testing.T) { t.Parallel() dto := taskDTO(&Task{ID: "7"}, "created") raw, err := json.Marshal(dto) if err != nil { t.Fatal(err) } var payload map[string]any if err := json.Unmarshal(raw, &payload); err != nil { t.Fatal(err) } for _, field := range []string{"source_task_ids", "llm_profile_ids"} { value, ok := payload[field].([]any) if !ok || len(value) != 0 { t.Fatalf("%s must serialize as [], payload=%s", field, raw) } } } func TestTaskLLMStateSnapshotIsConcurrentAndDetached(t *testing.T) { t.Parallel() task := &Task{ID: "7"} var wg sync.WaitGroup for writer := 0; writer < 2; writer++ { wg.Add(1) go func(offset int) { defer wg.Done() for i := 0; i < 500; i++ { id := int64(offset*1000 + i + 1) task.setLLMState(&id, &id, []int64{id, id + 1}, "ready", "") } }(writer) } for reader := 0; reader < 4; reader++ { wg.Add(1) go func() { defer wg.Done() for i := 0; i < 500; i++ { state := task.llmStateSnapshot() if len(state.ProfileIDs) > 0 { state.ProfileIDs[0] = -1 } _ = taskDTO(task, "running") } }() } wg.Wait() state := task.llmStateSnapshot() if len(state.ProfileIDs) != 2 || state.ProfileIDs[0] <= 0 { t.Fatalf("snapshot mutation leaked into task state: %+v", state) } } func TestTaskLLMStreamRetriesPreStreamQuotaOnNextProfile(t *testing.T) { t.Parallel() quota := errors.New("openai: status 402: insufficient_quota") first := &scriptedLLMProvider{err: quota} second := &scriptedLLMProvider{events: []llm.StreamEvent{{Type: llm.SETextDelta, Text: "fallback"}}} active := int64(11) transitions := 0 hooks := taskLLMStreamHooks{ current: func() (taskLLMSelection, error) { if active == 11 { return taskLLMSelection{profileID: 11, provider: first}, nil } return taskLLMSelection{profileID: 22, provider: second}, nil }, exhaust: func(selection taskLLMSelection, _ error) (db.TaskLLMTransition, error) { if selection.profileID != 11 { t.Fatalf("unexpected exhausted profile %d", selection.profileID) } active = 22 next := active return db.TaskLLMTransition{PreviousProfileID: 11, NextProfileID: &next, Advanced: true}, nil }, transition: func(taskLLMSelection, db.TaskLLMTransition, error) { transitions++ }, } events, err := collectTaskLLMStream(streamTaskLLM(context.Background(), "7", llm.CompletionRequest{}, hooks)) if err != nil { t.Fatal(err) } if len(events) != 1 || events[0].Text != "fallback" { t.Fatalf("unexpected fallback events: %+v", events) } if first.calls != 1 || second.calls != 1 || transitions != 1 { t.Fatalf("calls first=%d second=%d transitions=%d", first.calls, second.calls, transitions) } } func TestTaskLLMStreamRetriesAfterNonContentStartEvent(t *testing.T) { t.Parallel() quota := errors.New("anthropic: status 402: insufficient credit balance") first := &scriptedLLMProvider{ events: []llm.StreamEvent{{Type: llm.SEMessageStart}}, err: quota, } second := &scriptedLLMProvider{events: []llm.StreamEvent{ {Type: llm.SEMessageStart}, {Type: llm.SETextDelta, Text: "fallback"}, }} active := int64(11) hooks := taskLLMStreamHooks{ current: func() (taskLLMSelection, error) { if active == 11 { return taskLLMSelection{profileID: 11, provider: first}, nil } return taskLLMSelection{profileID: 22, provider: second}, nil }, exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) { active = 22 next := active return db.TaskLLMTransition{PreviousProfileID: 11, NextProfileID: &next, Advanced: true}, nil }, } events, err := collectTaskLLMStream(streamTaskLLM(context.Background(), "7", llm.CompletionRequest{}, hooks)) if err != nil { t.Fatal(err) } if len(events) != 2 || events[0].Type != llm.SEMessageStart || events[1].Text != "fallback" { t.Fatalf("non-content event from failed attempt leaked or fallback missing: %+v", events) } if first.calls != 1 || second.calls != 1 { t.Fatalf("calls first=%d second=%d", first.calls, second.calls) } } func TestTaskLLMStreamRetriesStalePreStreamFailureAgainstReplacementChain(t *testing.T) { t.Parallel() first := &scriptedLLMProvider{err: errors.New("status 402: insufficient_quota")} second := &scriptedLLMProvider{events: []llm.StreamEvent{{Type: llm.SETextDelta, Text: "replacement"}}} current := first hooks := taskLLMStreamHooks{ current: func() (taskLLMSelection, error) { profileID := int64(11) if current == second { profileID = 22 } return taskLLMSelection{profileID: profileID, provider: current}, nil }, exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) { current = second return db.TaskLLMTransition{PreviousProfileID: 11, Stale: true}, nil }, } events, err := collectTaskLLMStream(streamTaskLLM(context.Background(), "7", llm.CompletionRequest{}, hooks)) if err != nil || len(events) != 1 || events[0].Text != "replacement" { t.Fatalf("stale call did not continue on replacement chain: events=%+v err=%v", events, err) } } func TestTaskLLMStreamDoesNotSwitchForOrdinaryErrors(t *testing.T) { t.Parallel() tests := []string{ "openai: status 429: rate limit exceeded", "openai: status 401: invalid api key", "openai: status 500: internal server error", "dial tcp: network is unreachable", } for _, message := range tests { t.Run(message, func(t *testing.T) { provider := &scriptedLLMProvider{err: errors.New(message)} exhausted := 0 hooks := taskLLMStreamHooks{ current: func() (taskLLMSelection, error) { return taskLLMSelection{profileID: 11, provider: provider}, nil }, exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) { exhausted++ return db.TaskLLMTransition{}, nil }, } _, err := collectTaskLLMStream(streamTaskLLM(context.Background(), "7", llm.CompletionRequest{}, hooks)) if err == nil || err.Error() != message { t.Fatalf("error=%v, want %q", err, message) } if exhausted != 0 { t.Fatalf("ordinary error changed failover state %d times", exhausted) } }) } } func TestTaskLLMStreamAdvancesAfterMidStreamQuotaWithoutReplay(t *testing.T) { t.Parallel() quota := errors.New("anthropic: status 402: credit balance is too low") first := &scriptedLLMProvider{ events: []llm.StreamEvent{{Type: llm.SETextDelta, Text: "partial"}}, err: quota, } second := &scriptedLLMProvider{events: []llm.StreamEvent{{Type: llm.SETextDelta, Text: "next-call"}}} active := int64(11) hooks := taskLLMStreamHooks{ current: func() (taskLLMSelection, error) { if active == 11 { return taskLLMSelection{profileID: 11, provider: first}, nil } return taskLLMSelection{profileID: 22, provider: second}, nil }, exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) { active = 22 next := active return db.TaskLLMTransition{PreviousProfileID: 11, NextProfileID: &next, Advanced: true}, nil }, } events, err := collectTaskLLMStream(streamTaskLLM(context.Background(), "7", llm.CompletionRequest{}, hooks)) if len(events) != 1 || events[0].Text != "partial" || err == nil { t.Fatalf("mid-stream result events=%+v err=%v", events, err) } if first.calls != 1 || second.calls != 0 { t.Fatalf("mid-stream call was replayed: first=%d second=%d", first.calls, second.calls) } nextEvents, nextErr := collectTaskLLMStream(streamTaskLLM(context.Background(), "7", llm.CompletionRequest{}, hooks)) if nextErr != nil || len(nextEvents) != 1 || nextEvents[0].Text != "next-call" || second.calls != 1 { t.Fatalf("next call did not use fallback: events=%+v err=%v calls=%d", nextEvents, nextErr, second.calls) } } func TestTaskLLMStreamMarkFailureAfterToolStartSkipsWorkerReplay(t *testing.T) { t.Parallel() quota := errors.New("openai: status 402: insufficient_quota") markFailure := errors.New("database temporarily unavailable") provider := &scriptedLLMProvider{ events: []llm.StreamEvent{{Type: llm.SEToolUseStart}}, err: quota, } hooks := taskLLMStreamHooks{ current: func() (taskLLMSelection, error) { return taskLLMSelection{profileID: 11, provider: provider}, nil }, exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) { return db.TaskLLMTransition{}, markFailure }, } events, err := collectTaskLLMStream(streamTaskLLM(context.Background(), "7", llm.CompletionRequest{}, hooks)) if len(events) != 1 || events[0].Type != llm.SEToolUseStart || err == nil { t.Fatalf("committed tool stream events=%+v err=%v", events, err) } if !isTaskLLMRuntimeError(err) { t.Fatalf("committed mark failure must be router-classified: %T %v", err, err) } if retryableWorkerModelError(harness.ReasonModelError, err) { t.Fatal("worker must not replay after a committed tool stream when quota persistence fails") } if provider.calls != 1 { t.Fatalf("provider calls=%d, want 1", provider.calls) } } func TestTaskLLMRuntimeErrorSkipsWorkerReplay(t *testing.T) { partialStreamErr := &taskLLMError{ taskID: "7", chainExhausted: false, cause: errors.New("insufficient_quota after partial output"), } if retryableWorkerModelError(harness.ReasonModelError, partialStreamErr) { t.Fatal("worker must not replay an intent after the task router handled a mid-stream quota error") } if !retryableWorkerModelError(harness.ReasonModelError, errors.New("temporary network error")) { t.Fatal("ordinary model errors should retain the existing worker retry behavior") } } func TestEngineReadyForTaskOverride(t *testing.T) { t.Parallel() e := &Engine{} if e.Ready() { t.Fatal("engine without a global provider must not be globally ready") } e.SetAgentResolver(func(*Task) (*agent.Planner, *agent.Worker) { return &agent.Planner{}, &agent.Worker{} }) if !e.ReadyFor(&Task{ID: "7"}) { t.Fatal("task-specific provider chain must make that task ready") } } func TestResolvedTaskStatusKeepsExhaustedStartedTaskRunning(t *testing.T) { t.Parallel() engine := NewEngine(nil) engine.started.Store("7", true) task := &Task{ID: "7", Status: "running"} profileID := int64(11) task.setLLMState(&profileID, nil, []int64{profileID}, "chain_exhausted", "quota exhausted") s := &Server{engine: engine} if got := s.resolvedTaskStatus(task); got != "running" { t.Fatalf("exhausted started status=%q, want running", got) } task.Paused = true if got := s.resolvedTaskStatus(task); got != "paused" { t.Fatalf("paused exhausted status=%q, want paused", got) } } func TestAuthoritativeTaskResolverDoesNotFallBackWhenChainUnavailable(t *testing.T) { t.Parallel() e := &Engine{} e.UseLLM(&agent.Planner{}, &agent.Worker{}) if !e.ReadyFor(&Task{ID: "7"}) { t.Fatal("global provider should initially make task ready") } e.SetAuthoritativeAgentResolver(func(*Task) (*agent.Planner, *agent.Worker) { return nil, nil }) if e.ReadyFor(&Task{ID: "7"}) { t.Fatal("unavailable authoritative task chain must not fall back to global provider") } } func TestTaskLLMStreamStopsWhenChainExhausted(t *testing.T) { t.Parallel() provider := &scriptedLLMProvider{err: errors.New("billing_not_active")} transitions := 0 hooks := taskLLMStreamHooks{ current: func() (taskLLMSelection, error) { return taskLLMSelection{profileID: 11, provider: provider}, nil }, exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) { return db.TaskLLMTransition{PreviousProfileID: 11, Advanced: true, ChainExhausted: true}, nil }, transition: func(taskLLMSelection, db.TaskLLMTransition, error) { transitions++ }, } _, err := collectTaskLLMStream(streamTaskLLM(context.Background(), "7", llm.CompletionRequest{}, hooks)) if !isTaskLLMChainExhausted(err) { t.Fatalf("expected chain-exhausted error, got %v", err) } if provider.calls != 1 || transitions != 1 { t.Fatalf("calls=%d transitions=%d", provider.calls, transitions) } } func TestProfileDeleteRestoresQuotaBlockedIntentWhenFallbackAvailable(t *testing.T) { m, err := NewManager(t.TempDir(), "") if err != nil { t.Skipf("postgres unavailable (%v) - skipping", err) } defer m.Close() profileID, err := m.pg.SaveProfile(&db.LLMProfile{ Name: fmt.Sprintf("delete-quota-fallback-%d", time.Now().UnixNano()), Format: "openai", Model: "test-model", APIKey: "test-key", }) if err != nil { t.Fatal(err) } task, err := m.CreateTaskWithOptions("quota fallback", "resume intent", db.TaskCreateOptions{ LLMProfileIDs: []int64{profileID}, }) if err != nil { _ = m.pg.DeleteProfile(profileID) t.Fatal(err) } t.Cleanup(func() { _, _ = m.DeleteTask(task.ID, DeleteTaskOptions{}) _ = m.pg.DeleteProfile(profileID) }) intentID, err := task.Store.AddIntent(map[string]any{"action": "retry after fallback"}, 1, nil, "planner") if err != nil { t.Fatal(err) } taskID, err := parseTaskID(task.ID) if err != nil { t.Fatal(err) } if transition, err := m.pg.MarkTaskLLMProfileQuotaExhausted(taskID, profileID, "insufficient_quota"); err != nil || !transition.ChainExhausted { t.Fatalf("exhaust task chain: transition=%+v err=%v", transition, err) } if err := task.Store.SetIntentBlockedReason(intentID, db.IntentBlockedLLMQuota); err != nil { t.Fatal(err) } task.Paused = true if err := m.SetTaskPaused(task.ID, true); err != nil { t.Fatal(err) } // Removing the final explicit entry restores the legacy/global fallback. The // post-delete sync must reopen only quota-blocked work without unpausing it. if err := m.pg.DeleteProfile(profileID); err != nil { t.Fatal(err) } s := &Server{ m: m, llmOn: true, llmProv: &scriptedLLMProvider{}, provByProfile: map[int64]*provEntry{}, } s.restoreTasksAfterProfileDelete(m.pg) node, err := task.Store.GetNode(intentID) if err != nil { t.Fatal(err) } if node == nil || node.State != "open" || node.BlockedReason != "" { t.Fatalf("quota-blocked intent was not reopened: %+v", node) } if !task.Paused { t.Fatal("profile deletion must not resume a paused task") } state := task.llmStateSnapshot() if len(state.ProfileIDs) != 0 || state.FailoverState != "default" { t.Fatalf("task did not return to fallback state: %+v", state) } }