goclaw/internal/agent/loop_pipeline_callbacks_test.go

package agent

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

	"github.com/google/uuid"

	"github.com/nextlevelbuilder/goclaw/internal/bus"
	"github.com/nextlevelbuilder/goclaw/internal/pipeline"
	"github.com/nextlevelbuilder/goclaw/internal/providers"
	"github.com/nextlevelbuilder/goclaw/internal/tools"
	"github.com/nextlevelbuilder/goclaw/pkg/protocol"
)

type finalThinkingStreamProvider struct{}

func (p finalThinkingStreamProvider) Chat(context.Context, providers.ChatRequest) (*providers.ChatResponse, error) {
	return &providers.ChatResponse{Content: "final", Thinking: "non-stream thinking"}, nil
}

func (p finalThinkingStreamProvider) ChatStream(context.Context, providers.ChatRequest, func(providers.StreamChunk)) (*providers.ChatResponse, error) {
	return &providers.ChatResponse{Content: "final", Thinking: "final streamed thinking"}, nil
}

func (p finalThinkingStreamProvider) DefaultModel() string { return "test-model" }
func (p finalThinkingStreamProvider) Name() string         { return "test-provider" }

type requestCaptureProvider struct {
	request providers.ChatRequest
}

func (p *requestCaptureProvider) Chat(_ context.Context, req providers.ChatRequest) (*providers.ChatResponse, error) {
	p.request = req
	return &providers.ChatResponse{Content: "ok"}, nil
}

func (p *requestCaptureProvider) ChatStream(_ context.Context, req providers.ChatRequest, _ func(providers.StreamChunk)) (*providers.ChatResponse, error) {
	p.request = req
	return &providers.ChatResponse{Content: "ok"}, nil
}

func (p *requestCaptureProvider) DefaultModel() string { return "test-model" }
func (p *requestCaptureProvider) Name() string         { return "test-provider" }

type recordingSessionStore struct {
	*nopSessionStore
	added []providers.Message
}

func (s *recordingSessionStore) AddMessage(_ context.Context, _ string, msg providers.Message) {
	s.added = append(s.added, msg)
}

// retryOnceProvider fails the first Chat call with a transient 429 via RetryDo
// (which fires the injected retry hook), then succeeds — mimicking a flaky
// provider whose internal RetryDo consumes the run's retry hook.
type retryOnceProvider struct {
	calls int
	cfg   providers.RetryConfig
}

func (p *retryOnceProvider) Chat(ctx context.Context, _ providers.ChatRequest) (*providers.ChatResponse, error) {
	return providers.RetryDo(ctx, p.cfg, func() (*providers.ChatResponse, error) {
		p.calls++
		if p.calls == 1 {
			return nil, &providers.HTTPError{Status: 429, Body: "rate limited"}
		}
		return &providers.ChatResponse{Content: "ok"}, nil
	})
}

func (p *retryOnceProvider) ChatStream(ctx context.Context, _ providers.ChatRequest, _ func(providers.StreamChunk)) (*providers.ChatResponse, error) {
	return p.Chat(ctx, providers.ChatRequest{})
}

func (p *retryOnceProvider) DefaultModel() string { return "test-model" }
func (p *retryOnceProvider) Name() string         { return "test-provider" }

func TestMakeCallLLMEmitsRetryingOnTransientProviderError(t *testing.T) {
	col := &eventCollector{}
	loop := &Loop{id: "test-agent", onEvent: col.onEvent}
	req := &RunRequest{RunID: "run-1", SessionKey: "sess-1", Channel: "telegram"}
	state := &pipeline.RunState{
		Provider: &retryOnceProvider{
			cfg: providers.RetryConfig{Attempts: 2, MinDelay: time.Millisecond, MaxDelay: time.Millisecond},
		},
		Model:     "test-model",
		Iteration: 0,
	}

	resp, err := loop.makeCallLLM(req, col.onEvent)(context.Background(), state, providers.ChatRequest{})
	if err != nil {
		t.Fatalf("makeCallLLM returned error: %v", err)
	}
	if resp == nil || resp.Content != "ok" {
		t.Fatalf("unexpected response: %+v", resp)
	}

	retrying := col.filter(protocol.AgentEventRunRetrying)
	if len(retrying) != 1 {
		t.Fatalf("retrying events = %+v, want exactly one", retrying)
	}
	payload, ok := retrying[0].Payload.(map[string]string)
	if !ok || payload["attempt"] != "1" || payload["maxAttempts"] != "2" {
		t.Fatalf("retrying payload = %+v, want attempt=1 maxAttempts=2", retrying[0].Payload)
	}
}

func TestMakeCallLLM_StreamsFinalThinkingWhenNoThinkingChunkArrives(t *testing.T) {
	col := &eventCollector{}
	loop := &Loop{id: "test-agent", onEvent: col.onEvent}
	req := &RunRequest{
		RunID:      "run-1",
		SessionKey: "sess-1",
		Channel:    "telegram",
		Stream:     true,
	}
	state := &pipeline.RunState{
		Provider:  finalThinkingStreamProvider{},
		Model:     "test-model",
		Iteration: 0,
	}

	resp, err := loop.makeCallLLM(req, col.onEvent)(context.Background(), state, providers.ChatRequest{})
	if err != nil {
		t.Fatalf("makeCallLLM returned error: %v", err)
	}
	if resp == nil || resp.Thinking != "final streamed thinking" {
		t.Fatalf("stream response = %+v, want final thinking preserved", resp)
	}

	thinking := col.filter(protocol.ChatEventThinking)
	if len(thinking) != 1 {
		t.Fatalf("thinking events = %+v, want exactly one final thinking event", thinking)
	}
	payload, ok := thinking[0].Payload.(map[string]string)
	if !ok || payload["content"] != "final streamed thinking" {
		t.Fatalf("thinking payload = %+v", thinking[0].Payload)
	}
}

func TestMakeCallLLMPropagatesDelegationArtifactBridgeOptions(t *testing.T) {
	provider := &requestCaptureProvider{}
	loop := &Loop{id: "target-agent", agentUUID: uuid.New()}
	req := &RunRequest{RunID: "run-1", SessionKey: "delegate:session"}
	state := &pipeline.RunState{Provider: provider, Model: "test-model"}
	ctx := tools.WithToolWorkspace(context.Background(), "/runtime/outputs")
	ctx = tools.WithDelegationID(ctx, "delegation-id")
	ctx = tools.WithDelegationArtifactInputs(ctx, "/runtime/inputs")

	if _, err := loop.makeCallLLM(req, func(AgentEvent) {})(ctx, state, providers.ChatRequest{}); err != nil {
		t.Fatal(err)
	}
	if provider.request.Options[providers.OptDelegationID] != "delegation-id" ||
		provider.request.Options[providers.OptDelegationInputs] != "/runtime/inputs" {
		t.Fatalf("delegation options = %#v", provider.request.Options)
	}
}

func TestEnrichedInputMediaPersistsForNextTurn(t *testing.T) {
	sourceDir := t.TempDir()
	sourcePath := filepath.Join(sourceDir, "photo.png")
	if err := os.WriteFile(sourcePath, minimalPNG, 0o600); err != nil {
		t.Fatal(err)
	}
	workspace := t.TempDir()
	sessions := &recordingSessionStore{nopSessionStore: &nopSessionStore{}}
	loop := &Loop{sessions: sessions}
	req := &RunRequest{
		SessionKey: "session-media",
		Message:    `<media:image url="attachment://photo.png">`,
		Media: []bus.MediaFile{{
			Path:     sourcePath,
			MimeType: "image/png",
			Filename: "photo.png",
		}},
	}
	state := &pipeline.RunState{Messages: pipeline.NewMessageBuffer(providers.Message{Role: "system"})}
	state.Messages.SetHistory([]providers.Message{{Role: "user", Content: req.Message}})
	ctx := tools.WithToolWorkspace(context.Background(), workspace)

	if err := loop.makeEnrichMedia(req)(ctx, state); err != nil {
		t.Fatal(err)
	}
	if err := loop.makeFlushMessages(req)(ctx, req.SessionKey, nil); err != nil {
		t.Fatal(err)
	}

	if len(sessions.added) != 1 {
		t.Fatalf("persisted messages = %d, want 1", len(sessions.added))
	}
	persisted := sessions.added[0]
	if len(persisted.MediaRefs) != 1 {
		t.Fatalf("persisted MediaRefs = %#v, want one image ref", persisted.MediaRefs)
	}
	if strings.Contains(persisted.Content, workspace) {
		t.Fatalf("persisted content leaked workspace path: %q", persisted.Content)
	}
	if !strings.Contains(persisted.Content, `path=".uploads/`) {
		t.Fatalf("persisted content lacks logical image path: %q", persisted.Content)
	}

	nextTurnRefs := collectRefsByKind([]providers.Message{persisted}, nil, "image")
	if len(nextTurnRefs) != 1 || nextTurnRefs[0].ID != persisted.MediaRefs[0].ID {
		t.Fatalf("next-turn refs = %#v, want persisted exact ID", nextTurnRefs)
	}
}

func TestPromptCacheOptionsHelpers(t *testing.T) {
	tenantID := uuid.MustParse("11111111-1111-1111-1111-111111111111")
	agentID := uuid.MustParse("22222222-2222-2222-2222-222222222222")

	key1 := defaultPromptCacheKey(tenantID, agentID, "codex", "session-a")
	key2 := defaultPromptCacheKey(tenantID, agentID, "codex", "session-a")
	key3 := defaultPromptCacheKey(tenantID, agentID, "codex", "session-b")
	if key1 != key2 {
		t.Fatalf("defaultPromptCacheKey not stable: %q != %q", key1, key2)
	}
	if key1 == key3 {
		t.Fatal("defaultPromptCacheKey should vary by session")
	}
	if !strings.HasPrefix(key1, "goclaw/") {
		t.Fatalf("defaultPromptCacheKey = %q, want goclaw/ prefix", key1)
	}

	opts := map[string]any{}
	setDefaultPromptCacheOptions(opts, tenantID, agentID, "codex", "session-a")
	if opts[providers.OptPromptCacheKey] != key1 {
		t.Fatalf("prompt cache key = %v, want %s", opts[providers.OptPromptCacheKey], key1)
	}
	if opts[providers.OptPromptCacheRetention] != "24h" {
		t.Fatalf("prompt cache retention = %v, want 24h", opts[providers.OptPromptCacheRetention])
	}

	opts = map[string]any{
		providers.OptPromptCacheKey:       "custom-key",
		providers.OptPromptCacheRetention: "in_memory",
	}
	setDefaultPromptCacheOptions(opts, tenantID, agentID, "codex", "session-a")
	if opts[providers.OptPromptCacheKey] != "custom-key" || opts[providers.OptPromptCacheRetention] != "in_memory" {
		t.Fatalf("custom prompt cache options were overwritten: %+v", opts)
	}
}

func TestSupportsPromptCacheParams(t *testing.T) {
	if !supportsPromptCacheParams(providers.NewCodexProvider("codex", nil, "", "")) {
		t.Fatal("CodexProvider should support prompt cache params")
	}
	if supportsPromptCacheParams(finalThinkingStreamProvider{}) {
		t.Fatal("generic provider should not support prompt cache params")
	}
}

func TestResolveEffectiveContextWindow_UsesAgentConfigOnly(t *testing.T) {
	registry := &panicModelRegistry{}
	loop := &Loop{contextWindow: 128_000, modelRegistry: registry}
	if got := loop.resolveEffectiveContextWindow(); got != 128_000 {
		t.Fatalf("resolveEffectiveContextWindow() = %d, want 128000", got)
	}
}

type panicModelRegistry struct{}

func (*panicModelRegistry) Resolve(_, _ string) *providers.ModelSpec {
	panic("model registry must not participate in request budgeting")
}

func (*panicModelRegistry) Register(providers.ModelSpec) {
	panic("model registry must not participate in request budgeting")
}

func (*panicModelRegistry) Catalog(string) []providers.ModelSpec {
	panic("model registry must not participate in request budgeting")
}

func TestMakeUpdateMetadataStoresLastUsagePromptTokens(t *testing.T) {
	sessions := &nopSessionStore{}
	loop := &Loop{
		model:    "test-model",
		provider: finalThinkingStreamProvider{},
		sessions: sessions,
	}
	req := &RunRequest{Channel: "telegram"}

	update := loop.makeUpdateMetadata(req)
	err := update(context.Background(), "sess-1",
		providers.Usage{PromptTokens: 225000, CompletionTokens: 3000},
		providers.Usage{PromptTokens: 70000, CompletionTokens: 100},
		12,
	)
	if err != nil {
		t.Fatalf("update metadata error: %v", err)
	}
	if sessions.inputTokens != 225000 || sessions.outputTokens != 3000 {
		t.Fatalf("accumulated tokens = %d/%d, want total run 225000/3000", sessions.inputTokens, sessions.outputTokens)
	}
	// Upstream 503909d3 calibration: SetLastPromptTokens stores the final call's
	// full context size (Usage.ContextTokens(), which adds cached segments back
	// for Anthropic-style usage) PLUS the final completion — the reply joins
	// history so it occupies the next request's prompt. No cache tokens here, so
	// ContextTokens()=70000; +100 completion = 70100.
	if sessions.setLastTokens != 70100 || sessions.setLastMsgCount != 12 {
		t.Fatalf("last prompt calibration = %d/%d, want last request 70100/12", sessions.setLastTokens, sessions.setLastMsgCount)
	}
}

// A Function-nil tool definition (e.g. the native image_generation sentinel,
// providers.ToolDefinition{Type: "image_generation"}) must not panic the
// mcp-def counter. Regression for the v3.14.0 nil-pointer crash.
func TestCountMCPToolDefs_SkipsNilFunction(t *testing.T) {
	defs := []providers.ToolDefinition{
		{Type: "image_generation"}, // Function == nil
		{Function: &providers.ToolFunctionSchema{Name: "mcp_notion_search"}},
		{Function: &providers.ToolFunctionSchema{Name: " mcp_slack_post "}},
		{Function: &providers.ToolFunctionSchema{Name: "read_file"}},
	}

	if got := countMCPToolDefs(defs); got != 2 {
		t.Errorf("countMCPToolDefs = %d, want 2", got)
	}
}

// The image_generation sentinel must carry a non-nil Function so the many
// pipeline/provider sites that read td.Function.Name (think_stage, codex_build,
// shouldRetryTaskMCP, history tool names, …) never nil-deref. Root-cause guard
// for the v3.14.0 crash — one landmine removed instead of guarding every site.
func TestImageGenToolDef_FunctionNonNil(t *testing.T) {
	if imageGenToolDef.Type != "image_generation" {
		t.Fatalf("sentinel Type = %q, want image_generation", imageGenToolDef.Type)
	}
	if imageGenToolDef.Function == nil {
		t.Fatal("sentinel Function must be non-nil to avoid downstream nil-deref")
	}
	if imageGenToolDef.Function.Name != "image_generation" {
		t.Errorf("sentinel Function.Name = %q, want image_generation", imageGenToolDef.Function.Name)
	}
}