diff --git a/server/internal/daemon/daemon.go b/server/internal/daemon/daemon.go index 58608cf50..a25340698 100644 --- a/server/internal/daemon/daemon.go +++ b/server/internal/daemon/daemon.go @@ -988,12 +988,19 @@ func (d *Daemon) runTask(ctx context.Context, task Task, provider string, taskLo return TaskResult{}, err } - // Fallback: if session resume failed, retry with a fresh session. - if result.Status == "failed" && task.PriorSessionID != "" { + // Fallback: if session resume failed before establishing a session, retry + // with a fresh session. We check SessionID == "" to distinguish a resume + // failure (no session established) from a failure during actual execution. + if result.Status == "failed" && task.PriorSessionID != "" && result.SessionID == "" { + firstUsage := result.Usage taskLog.Warn("session resume failed, retrying with fresh session", "error", result.Error) execOpts.ResumeSessionID = "" - if retryResult, retryTools, retryErr := d.executeAndDrain(ctx, backend, prompt, execOpts, taskLog, task.ID); retryErr == nil { + retryResult, retryTools, retryErr := d.executeAndDrain(ctx, backend, prompt, execOpts, taskLog, task.ID) + if retryErr != nil { + taskLog.Error("fresh session also failed to start", "error", retryErr) + } else { result = retryResult + result.Usage = mergeUsage(firstUsage, result.Usage) tools = retryTools } } @@ -1182,6 +1189,28 @@ func (d *Daemon) executeAndDrain(ctx context.Context, backend agent.Backend, pro return result, toolCount.Load(), nil } +func mergeUsage(a, b map[string]agent.TokenUsage) map[string]agent.TokenUsage { + if len(a) == 0 { + return b + } + if len(b) == 0 { + return a + } + merged := make(map[string]agent.TokenUsage, len(a)+len(b)) + for model, u := range a { + merged[model] = u + } + for model, u := range b { + existing := merged[model] + existing.InputTokens += u.InputTokens + existing.OutputTokens += u.OutputTokens + existing.CacheReadTokens += u.CacheReadTokens + existing.CacheWriteTokens += u.CacheWriteTokens + merged[model] = existing + } + return merged +} + // repoDataToInfo converts daemon RepoData to repocache RepoInfo. func repoDataToInfo(repos []RepoData) []repocache.RepoInfo { info := make([]repocache.RepoInfo, len(repos)) diff --git a/server/internal/daemon/daemon_test.go b/server/internal/daemon/daemon_test.go index 868970d61..0f56b9fbe 100644 --- a/server/internal/daemon/daemon_test.go +++ b/server/internal/daemon/daemon_test.go @@ -1,9 +1,15 @@ package daemon import ( + "context" + "log/slog" "net/http" + "net/http/httptest" "strings" + "sync/atomic" "testing" + + "github.com/multica-ai/multica/server/pkg/agent" ) func TestNormalizeServerBaseURL(t *testing.T) { @@ -83,3 +89,147 @@ func TestIsWorkspaceNotFoundError(t *testing.T) { t.Fatal("did not expect 500 to be treated as workspace not found") } } + +func TestMergeUsage(t *testing.T) { + t.Parallel() + + a := map[string]agent.TokenUsage{ + "model-a": {InputTokens: 10, OutputTokens: 5}, + } + b := map[string]agent.TokenUsage{ + "model-a": {InputTokens: 20, OutputTokens: 10, CacheReadTokens: 3}, + "model-b": {InputTokens: 100}, + } + merged := mergeUsage(a, b) + + if got := merged["model-a"]; got.InputTokens != 30 || got.OutputTokens != 15 || got.CacheReadTokens != 3 { + t.Fatalf("model-a: expected {30,15,3,0}, got %+v", got) + } + if got := merged["model-b"]; got.InputTokens != 100 { + t.Fatalf("model-b: expected InputTokens=100, got %+v", got) + } + + if got := mergeUsage(nil, b); len(got) != 2 { + t.Fatal("mergeUsage(nil, b) should return b") + } + if got := mergeUsage(a, nil); len(got) != 1 { + t.Fatal("mergeUsage(a, nil) should return a") + } +} + +// fakeBackend is a test double for agent.Backend that returns preconfigured +// results. Each call to Execute pops the next entry from the results slice. +type fakeBackend struct { + calls []agent.ExecOptions + results []agent.Result + errors []error + idx atomic.Int32 +} + +func (b *fakeBackend) Execute(_ context.Context, _ string, opts agent.ExecOptions) (*agent.Session, error) { + i := int(b.idx.Add(1)) - 1 + b.calls = append(b.calls, opts) + if i < len(b.errors) && b.errors[i] != nil { + return nil, b.errors[i] + } + msgCh := make(chan agent.Message) + resCh := make(chan agent.Result, 1) + close(msgCh) + resCh <- b.results[i] + return &agent.Session{Messages: msgCh, Result: resCh}, nil +} + +func newTestDaemon(t *testing.T) *Daemon { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(srv.Close) + return &Daemon{ + client: NewClient(srv.URL), + logger: slog.Default(), + } +} + +func TestExecuteAndDrain_ResumeFailureFallback(t *testing.T) { + t.Parallel() + + d := newTestDaemon(t) + ctx := context.Background() + taskLog := slog.Default() + + fb := &fakeBackend{ + results: []agent.Result{ + {Status: "failed", Error: "session not found", Usage: map[string]agent.TokenUsage{ + "m1": {InputTokens: 5}, + }}, + {Status: "completed", Output: "done", SessionID: "new-sess", Usage: map[string]agent.TokenUsage{ + "m1": {InputTokens: 10, OutputTokens: 20}, + }}, + }, + } + + // First attempt: resume fails (no SessionID in result). + opts := agent.ExecOptions{ResumeSessionID: "stale-id"} + result, _, err := d.executeAndDrain(ctx, fb, "prompt", opts, taskLog, "task-1") + if err != nil { + t.Fatalf("first call error: %v", err) + } + if result.Status != "failed" || result.SessionID != "" { + t.Fatalf("expected failed result with empty SessionID, got %+v", result) + } + + // Simulate the retry logic from runTask. + if result.Status == "failed" && result.SessionID == "" { + firstUsage := result.Usage + opts.ResumeSessionID = "" + retryResult, _, retryErr := d.executeAndDrain(ctx, fb, "prompt", opts, taskLog, "task-1") + if retryErr != nil { + t.Fatalf("retry error: %v", retryErr) + } + result = retryResult + result.Usage = mergeUsage(firstUsage, result.Usage) + } + + if result.Status != "completed" || result.Output != "done" { + t.Fatalf("expected completed result, got %+v", result) + } + if result.SessionID != "new-sess" { + t.Fatalf("expected new-sess, got %s", result.SessionID) + } + // Usage should be merged. + if u := result.Usage["m1"]; u.InputTokens != 15 || u.OutputTokens != 20 { + t.Fatalf("expected merged usage {15,20}, got %+v", u) + } + // Second call should NOT have ResumeSessionID. + if fb.calls[1].ResumeSessionID != "" { + t.Fatal("retry should not have ResumeSessionID") + } +} + +func TestExecuteAndDrain_NoRetryWhenSessionEstablished(t *testing.T) { + t.Parallel() + + d := newTestDaemon(t) + + fb := &fakeBackend{ + results: []agent.Result{ + {Status: "failed", Error: "model error", SessionID: "valid-sess"}, + }, + } + + opts := agent.ExecOptions{ResumeSessionID: "some-id"} + result, _, err := d.executeAndDrain(context.Background(), fb, "p", opts, slog.Default(), "t") + if err != nil { + t.Fatal(err) + } + + // SessionID is set → session was established → should NOT retry. + shouldRetry := result.Status == "failed" && result.SessionID == "" + if shouldRetry { + t.Fatal("should not retry when SessionID is present") + } + if int(fb.idx.Load()) != 1 { + t.Fatalf("expected 1 call, got %d", fb.idx.Load()) + } +}