agentu/internal/session/manager_test.go

package session

import (
	"context"
	"path/filepath"
	"strings"
	"testing"

	"agentu/internal/agent"
	"agentu/pkg/config"
	"agentu/pkg/llm"
)

type recordingProvider struct {
	req llm.ChatRequest
}

func (p *recordingProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
	p.req = req
	return emit(llm.StreamEvent{Content: "ok"})
}

func testConfig() *config.Config {
	return &config.Config{
		Providers: map[string]config.ProviderConfig{
			"one": {
				Name:               "one",
				BaseURL:            "https://one.example.com",
				APIKey:             "sk-test",
				Model:              "model-a",
				Models:             []string{"model-a", "model-b"},
				Thinking:           "none",
				ThinkingParam:      "thinking",
				ThinkingConfigured: true,
			},
		},
	}
}

func testManager(t *testing.T) (*Manager, *agent.Agent, *recordingProvider) {
	t.Helper()
	cfg := testConfig()
	provider := &recordingProvider{}
	assistant := agent.New(agent.Options{
		Provider:        provider,
		ProviderName:    "one",
		Model:           "model-a",
		Thinking:        "none",
		ThinkingParam:   "thinking",
		ThinkingEnabled: true,
	})
	store, err := NewStore(filepath.Join(t.TempDir(), "sessions"))
	if err != nil {
		t.Fatal(err)
	}
	manager, err := NewManager(cfg, assistant, nil, store)
	if err != nil {
		t.Fatal(err)
	}
	return manager, assistant, provider
}

func TestManagerSwitchesModelAndThinking(t *testing.T) {
	manager, assistant, provider := testManager(t)

	if _, err := manager.SetModel("model-b"); err != nil {
		t.Fatal(err)
	}
	if _, err := manager.SetThinking("high"); err != nil {
		t.Fatal(err)
	}
	var out strings.Builder
	if err := assistant.RunTurn(context.Background(), "hi", &out, nil); err != nil {
		t.Fatal(err)
	}

	if provider.req.Model != "model-b" {
		t.Fatalf("model = %q", provider.req.Model)
	}
	if provider.req.Extra["thinking"] != "high" {
		t.Fatalf("thinking = %#v", provider.req.Extra["thinking"])
	}
}

func TestManagerRejectsUnknownModelAndThinking(t *testing.T) {
	cfg := &config.Config{
		Providers: map[string]config.ProviderConfig{
			"one": {
				Name:    "one",
				BaseURL: "https://one.example.com",
				APIKey:  "sk-test",
				Model:   "model-a",
				Models:  []string{"model-a"},
			},
		},
	}
	assistant := agent.New(agent.Options{Provider: &recordingProvider{}, Model: "model-a"})
	store, _ := NewStore(filepath.Join(t.TempDir(), "sessions"))
	manager, err := NewManager(cfg, assistant, nil, store)
	if err != nil {
		t.Fatal(err)
	}

	if _, err := manager.SetModel("model-b"); err == nil {
		t.Fatal("expected unknown model error")
	}
	if _, err := manager.SetThinking("huge"); err == nil {
		t.Fatal("expected invalid thinking error")
	}
}

func TestManagerExposesModelCommandCandidates(t *testing.T) {
	manager, _, _ := testManager(t)
	models := manager.AvailableModels()
	if len(models) != 2 || models[0] != "model-a" || models[1] != "model-b" {
		t.Fatalf("AvailableModels() = %v", models)
	}
	models[0] = "mutated"
	if got := manager.AvailableModels()[0]; got != "model-a" {
		t.Fatalf("mutation leaked: %q", got)
	}
}

func TestManagerSaveAndResume(t *testing.T) {
	manager, assistant, _ := testManager(t)

	// Start new session and have a conversation
	manager.newSession()
	var out strings.Builder
	assistant.RunTurn(context.Background(), "hello world", &out, nil)
	assistant.SetLastUsage(&llm.Usage{PromptTokens: 10, CompletionTokens: 5, TotalTokens: 15})

	// Save
	if err := manager.Save(); err != nil {
		t.Fatalf("save: %v", err)
	}
	savedID := manager.CurrentSessionID()
	if savedID == "" {
		t.Fatal("expected session ID after save")
	}
	// Name should be auto-derived from first user message
	if manager.CurrentSessionName() != "hello world" {
		t.Fatalf("session name = %q, want %q", manager.CurrentSessionName(), "hello world")
	}

	// Start a new session
	manager.newSession()
	if manager.CurrentSessionID() == savedID {
		t.Fatal("new session should have different ID")
	}

	// Resume
	if err := manager.Resume(savedID); err != nil {
		t.Fatalf("resume: %v", err)
	}
	if manager.CurrentSessionID() != savedID {
		t.Fatalf("after resume, ID = %q, want %q", manager.CurrentSessionID(), savedID)
	}
	msgs := assistant.Messages()
	hasUser := false
	for _, m := range msgs {
		if m.Role == llm.RoleUser && m.Content == "hello world" {
			hasUser = true
		}
	}
	if !hasUser {
		t.Fatal("resumed session should contain original user message")
	}
	if usage := assistant.LastUsage(); usage == nil || usage.TotalTokens != 15 {
		t.Fatalf("resumed usage = %#v", usage)
	}
}

func TestManagerRename(t *testing.T) {
	manager, _, _ := testManager(t)
	manager.newSession()

	msg, err := manager.Rename("my cool session")
	if err != nil {
		t.Fatalf("rename: %v", err)
	}
	if !strings.Contains(msg, "my cool session") {
		t.Fatalf("rename message = %q", msg)
	}
	if manager.CurrentSessionName() != "my cool session" {
		t.Fatalf("name = %q", manager.CurrentSessionName())
	}

	// Verify persisted
	_, err = manager.store.Load(manager.CurrentSessionID())
	if err != nil {
		t.Fatalf("load after rename: %v", err)
	}
}

func TestManagerRenameEmpty(t *testing.T) {
	manager, _, _ := testManager(t)
	manager.newSession()
	_, err := manager.Rename("  ")
	if err == nil {
		t.Fatal("expected error for empty name")
	}
}

func TestManagerSessions(t *testing.T) {
	manager, _, _ := testManager(t)

	// No sessions yet
	out, err := manager.Sessions()
	if err != nil {
		t.Fatal(err)
	}
	if !strings.Contains(out, "No saved sessions") {
		t.Fatalf("expected 'no saved sessions', got: %s", out)
	}

	// Create and save a session
	manager.newSession()
	manager.Save()

	out, err = manager.Sessions()
	if err != nil {
		t.Fatal(err)
	}
	if strings.Contains(out, "No saved sessions") {
		t.Fatal("expected session list")
	}
}

func TestManagerNewSession(t *testing.T) {
	manager, assistant, _ := testManager(t)
	manager.newSession()
	var out strings.Builder
	assistant.RunTurn(context.Background(), "first message", &out, nil)

	msg, err := manager.NewSession()
	if err != nil {
		t.Fatalf("new session: %v", err)
	}
	if !strings.Contains(msg, "Starting new session") {
		t.Fatalf("message = %q", msg)
	}
	// Old session should be saved
	metas, _ := manager.store.List()
	if len(metas) != 1 {
		t.Fatalf("expected 1 saved session, got %d", len(metas))
	}
}

func TestManagerResumeLatest(t *testing.T) {
	manager, assistant, _ := testManager(t)

	// Create two sessions
	manager.newSession()
	var out strings.Builder
	assistant.RunTurn(context.Background(), "session one", &out, nil)
	manager.Save()
	id1 := manager.CurrentSessionID()

	manager.NewSession()
	assistant.RunTurn(context.Background(), "session two", &out, nil)
	manager.Save()
	id2 := manager.CurrentSessionID()

	// Clear and resume latest
	manager.newSession()
	if err := manager.ResumeLatest(); err != nil {
		t.Fatalf("resume latest: %v", err)
	}
	// Should resume the most recently updated (id2)
	if manager.CurrentSessionID() != id2 {
		t.Fatalf("expected to resume %s, got %s", id2, manager.CurrentSessionID())
	}

	// Also verify id1 is resumable
	if err := manager.Resume(id1); err != nil {
		t.Fatalf("resume id1: %v", err)
	}
	if manager.CurrentSessionID() != id1 {
		t.Fatalf("expected %s, got %s", id1, manager.CurrentSessionID())
	}
}

type compactingProvider struct {
	summaryCalls int
	answerCalls  int
}

func (p *compactingProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
	if len(req.Messages) == 2 && req.Messages[0].Role == llm.RoleSystem {
		p.summaryCalls++
		return emit(llm.StreamEvent{Content: "summary"})
	}
	p.answerCalls++
	return emit(llm.StreamEvent{Content: "answer"})
}

func TestManagerUsesModelContextForAgentCompaction(t *testing.T) {
	provider := &compactingProvider{}
	cfg := &config.Config{Providers: map[string]config.ProviderConfig{
		"one": {
			Name:          "one",
			BaseURL:       "https://one.example.com",
			APIKey:        "sk-test",
			Model:         "model-a",
			Models:        []string{"model-a"},
			ModelConfigs:  map[string]config.ModelConfig{"model-a": {ID: "model-a", Thinking: "none", ContextTokens: 1}},
			ContextTokens: 1,
		},
	}}
	assistant := agent.New(agent.Options{Provider: provider, ProviderName: "one", Model: "model-a"})
	store, err := NewStore(filepath.Join(t.TempDir(), "sessions"))
	if err != nil {
		t.Fatal(err)
	}
	manager, err := NewManager(cfg, assistant, nil, store)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := manager.SetModel("model-a"); err != nil {
		t.Fatal(err)
	}
	assistant.SetMessages([]llm.Message{
		{Role: llm.RoleSystem, Content: "system"},
		{Role: llm.RoleUser, Content: "old question 1"},
		{Role: llm.RoleAssistant, Content: "old answer 1"},
		{Role: llm.RoleUser, Content: "old question 2"},
		{Role: llm.RoleAssistant, Content: "old answer 2"},
		{Role: llm.RoleUser, Content: "recent question 1"},
		{Role: llm.RoleAssistant, Content: "recent answer 1"},
		{Role: llm.RoleUser, Content: "recent question 2"},
		{Role: llm.RoleAssistant, Content: "recent answer 2"},
	})

	var out strings.Builder
	if err := assistant.RunTurn(context.Background(), "new question", &out, nil); err != nil {
		t.Fatal(err)
	}
	if provider.summaryCalls != 1 {
		t.Fatalf("summaryCalls = %d, want 1", provider.summaryCalls)
	}
	if provider.answerCalls != 1 {
		t.Fatalf("answerCalls = %d, want 1", provider.answerCalls)
	}
}