/*
 * Copyright 2025 CloudWeGo Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package adk

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"sync"
	"sync/atomic"
	"testing"

	"github.com/stretchr/testify/assert"

	"github.com/cloudwego/eino/components/model"
	"github.com/cloudwego/eino/components/tool"
	"github.com/cloudwego/eino/compose"
	"github.com/cloudwego/eino/schema"
)

type interruptTestToolsHandler struct {
	*BaseChatModelAgentMiddleware
	tools []tool.BaseTool
}

func TestPreprocessADKCheckpoint(t *testing.T) {
	t.Run("no-op when missing markers", func(t *testing.T) {
		in := []byte("random")
		out := preprocessADKCheckpoint(append([]byte(nil), in...))
		assert.Equal(t, in, out)
	})

	t.Run("rewrite legacy name for v0.8.0-v0.8.3", func(t *testing.T) {
		const (
			lenPrefixedReactStateName         = "\x15" + stateGobNameV07
			lenPrefixedCompatName             = "\x15" + stateGobNameV080
			lenPrefixedStateSerializationName = "\x12stateSerialization"
		)

		in := []byte(lenPrefixedReactStateName + "xxx" + lenPrefixedStateSerializationName + "yyy")
		out := preprocessADKCheckpoint(append([]byte(nil), in...))
		assert.True(t, bytes.Contains(out, []byte(lenPrefixedCompatName)))
		assert.False(t, bytes.Contains(out, []byte(lenPrefixedReactStateName)))
	})
}

func (h *interruptTestToolsHandler) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) {
	runCtx.Tools = append(runCtx.Tools, h.tools...)
	return ctx, runCtx, nil
}

func TestSaveAgentEventWrapper(t *testing.T) {
	sr, sw := schema.Pipe[Message](1)
	sw.Send(schema.UserMessage("test"), nil)
	sw.Close()
	sr = sr.Copy(2)[1]

	w := &agentEventWrapper{
		AgentEvent: &AgentEvent{
			Output: &AgentOutput{
				MessageOutput: &MessageVariant{
					IsStreaming:   true,
					MessageStream: sr,
				},
			},
			RunPath: []RunStep{
				{
					"a1",
				},
				{
					"a2",
				},
			},
		},
		mu:                  sync.Mutex{},
		concatenatedMessage: nil,
	}

	_, err := getMessageFromWrappedEvent(w)
	assert.NoError(t, err)

	buf, err := w.GobEncode()
	assert.NoError(t, err)
	assert.NoError(t, err)

	w1 := &agentEventWrapper{}
	err = w1.GobDecode(buf)
	assert.NoError(t, err)
}

func TestInterruptFunctionsPopulateInterruptContextsImmediately(t *testing.T) {
	ctx := context.Background()
	ctx, _ = initRunCtx(ctx, "TestAgent", &AgentInput{Messages: []Message{}})
	ctx = AppendAddressSegment(ctx, AddressSegmentAgent, "TestAgent")

	t.Run("Interrupt populates InterruptContexts", func(t *testing.T) {
		event := Interrupt(ctx, "test info")
		assert.NotNil(t, event.Action)
		assert.NotNil(t, event.Action.Interrupted)
		assert.NotNil(t, event.Action.Interrupted.InterruptContexts)
		assert.Equal(t, 1, len(event.Action.Interrupted.InterruptContexts))
		assert.Equal(t, "test info", event.Action.Interrupted.InterruptContexts[0].Info)
		assert.True(t, event.Action.Interrupted.InterruptContexts[0].IsRootCause)
		assert.Equal(t, Address{
			{Type: AddressSegmentAgent, ID: "TestAgent"},
		}, event.Action.Interrupted.InterruptContexts[0].Address)
	})

	t.Run("StatefulInterrupt populates InterruptContexts", func(t *testing.T) {
		event := StatefulInterrupt(ctx, "stateful info", "my state")
		assert.NotNil(t, event.Action)
		assert.NotNil(t, event.Action.Interrupted)
		assert.NotNil(t, event.Action.Interrupted.InterruptContexts)
		assert.Equal(t, 1, len(event.Action.Interrupted.InterruptContexts))
		assert.Equal(t, "stateful info", event.Action.Interrupted.InterruptContexts[0].Info)
		assert.True(t, event.Action.Interrupted.InterruptContexts[0].IsRootCause)
	})

	t.Run("CompositeInterrupt populates InterruptContexts with filtered parent chain", func(t *testing.T) {
		subCtx := AppendAddressSegment(ctx, AddressSegmentAgent, "SubAgent")
		subEvent := Interrupt(subCtx, "sub info")
		event := CompositeInterrupt(ctx, "composite info", "composite state", subEvent.Action.internalInterrupted)
		assert.NotNil(t, event.Action)
		assert.NotNil(t, event.Action.Interrupted)
		assert.NotNil(t, event.Action.Interrupted.InterruptContexts)
		assert.Equal(t, 1, len(event.Action.Interrupted.InterruptContexts))

		rootCause := event.Action.Interrupted.InterruptContexts[0]
		assert.Equal(t, "sub info", rootCause.Info)
		assert.True(t, rootCause.IsRootCause)
		assert.Equal(t, Address{
			{Type: AddressSegmentAgent, ID: "TestAgent"},
			{Type: AddressSegmentAgent, ID: "SubAgent"},
		}, rootCause.Address)

		assert.NotNil(t, rootCause.Parent, "Parent should not be nil for composite interrupt")
		assert.Equal(t, "composite info", rootCause.Parent.Info)
		assert.Equal(t, Address{
			{Type: AddressSegmentAgent, ID: "TestAgent"},
		}, rootCause.Parent.Address)
	})

	t.Run("Address only contains agent/tool segments", func(t *testing.T) {
		event := Interrupt(ctx, "test info")
		addr := event.Action.Interrupted.InterruptContexts[0].Address
		for _, seg := range addr {
			assert.True(t, seg.Type == AddressSegmentAgent || seg.Type == AddressSegmentTool,
				"Address should only contain agent/tool segments, got: %s", seg.Type)
		}
	})
}

func TestSimpleInterrupt(t *testing.T) {
	data := "hello world"
	agent := &myAgent{
		runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
			iter, generator := NewAsyncIteratorPair[*AgentEvent]()
			generator.Send(&AgentEvent{
				Output: &AgentOutput{
					MessageOutput: &MessageVariant{
						IsStreaming: true,
						Message:     nil,
						MessageStream: schema.StreamReaderFromArray([]Message{
							schema.UserMessage("hello "),
							schema.UserMessage("world"),
						}),
					},
				},
			})
			intEvent := Interrupt(ctx, data)
			intEvent.Action.Interrupted.Data = data
			generator.Send(intEvent)
			generator.Close()
			return iter
		},
		resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {
			assert.True(t, info.WasInterrupted)
			assert.Nil(t, info.InterruptState)
			assert.True(t, info.EnableStreaming)
			assert.Equal(t, data, info.Data)

			assert.True(t, info.IsResumeTarget)
			iter, generator := NewAsyncIteratorPair[*AgentEvent]()
			generator.Close()
			return iter
		},
	}
	store := newMyStore()
	ctx := context.Background()
	runner := NewRunner(ctx, RunnerConfig{
		Agent:           agent,
		EnableStreaming: true,
		CheckPointStore: store,
	})
	iter := runner.Query(ctx, "hello world", WithCheckPointID("1"))
	_, ok := iter.Next()
	assert.True(t, ok)
	interruptEvent, ok := iter.Next()
	assert.True(t, ok)
	assert.Equal(t, data, interruptEvent.Action.Interrupted.Data)
	assert.NotEmpty(t, interruptEvent.Action.Interrupted.InterruptContexts[0].ID)
	assert.True(t, interruptEvent.Action.Interrupted.InterruptContexts[0].IsRootCause)
	assert.Equal(t, data, interruptEvent.Action.Interrupted.InterruptContexts[0].Info)
	assert.Equal(t, Address{{Type: AddressSegmentAgent, ID: "myAgent"}},
		interruptEvent.Action.Interrupted.InterruptContexts[0].Address)
	_, ok = iter.Next()
	assert.False(t, ok)

	iter, err := runner.ResumeWithParams(ctx, "1", &ResumeParams{
		Targets: map[string]any{
			interruptEvent.Action.Interrupted.InterruptContexts[0].ID: nil,
		},
	})
	assert.NoError(t, err)
	_, ok = iter.Next()
	assert.False(t, ok)
}

func TestMultiAgentInterrupt(t *testing.T) {
	ctx := context.Background()
	sa1 := &myAgent{
		name: "sa1",
		runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
			iter, generator := NewAsyncIteratorPair[*AgentEvent]()
			generator.Send(&AgentEvent{
				AgentName: "sa1",
				Action: &AgentAction{
					TransferToAgent: &TransferToAgentAction{
						DestAgentName: "sa2",
					},
				},
			})
			generator.Close()
			return iter
		},
	}
	sa2 := &myAgent{
		name: "sa2",
		runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
			iter, generator := NewAsyncIteratorPair[*AgentEvent]()
			intEvent := StatefulInterrupt(ctx, "hello world", "temp state")
			intEvent.Action.Interrupted.Data = "hello world"
			generator.Send(intEvent)
			generator.Close()
			return iter
		},
		resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {
			assert.NotNil(t, info)
			assert.Equal(t, info.Data, "hello world")

			assert.True(t, info.WasInterrupted)
			assert.NotNil(t, info.InterruptState)
			assert.Equal(t, "temp state", info.InterruptState)

			assert.True(t, info.IsResumeTarget)
			assert.NotNil(t, info.ResumeData)
			assert.Equal(t, "resume data", info.ResumeData)

			iter, generator := NewAsyncIteratorPair[*AgentEvent]()
			generator.Send(&AgentEvent{
				AgentName: "sa2",
				Output: &AgentOutput{
					MessageOutput: &MessageVariant{Message: schema.UserMessage(info.ResumeData.(string))},
				},
			})
			generator.Close()
			return iter
		},
	}
	a, err := SetSubAgents(ctx, sa1, []Agent{sa2})
	assert.NoError(t, err)
	runner := NewRunner(ctx, RunnerConfig{
		Agent:           a,
		EnableStreaming: false,
		CheckPointStore: newMyStore(),
	})
	iter := runner.Query(ctx, "", WithCheckPointID("1"))
	event, ok := iter.Next()
	assert.True(t, ok)
	assert.NotNil(t, event.Action.TransferToAgent)
	event, ok = iter.Next()
	assert.True(t, ok)
	assert.NotNil(t, event.Action.Interrupted)
	assert.Equal(t, 1, len(event.Action.Interrupted.InterruptContexts))
	assert.Equal(t, "hello world", event.Action.Interrupted.InterruptContexts[0].Info)
	assert.True(t, event.Action.Interrupted.InterruptContexts[0].IsRootCause)
	assert.Equal(t, Address{
		{Type: AddressSegmentAgent, ID: "sa1"},
		{Type: AddressSegmentAgent, ID: "sa2"},
	}, event.Action.Interrupted.InterruptContexts[0].Address)
	assert.NotEmpty(t, event.Action.Interrupted.InterruptContexts[0].ID)

	interruptID := event.Action.Interrupted.InterruptContexts[0].ID
	_, ok = iter.Next()
	assert.False(t, ok)

	iter, err = runner.ResumeWithParams(ctx, "1", &ResumeParams{
		Targets: map[string]any{
			interruptID: "resume data",
		},
	})
	assert.NoError(t, err)
	event, ok = iter.Next()
	assert.True(t, ok)
	assert.Equal(t, event.Output.MessageOutput.Message.Content, "resume data")
	_, ok = iter.Next()
	assert.False(t, ok)
}

func TestWorkflowInterrupt(t *testing.T) {
	ctx := context.Background()
	sa1 := &myAgent{
		name: "sa1",
		runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
			iter, generator := NewAsyncIteratorPair[*AgentEvent]()

			intEvent := Interrupt(ctx, "sa1 interrupt data")
			intEvent.Action.Interrupted.Data = "sa1 interrupt data"
			generator.Send(intEvent)
			generator.Close()
			return iter
		},
		resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {
			assert.Equal(t, info.InterruptInfo.Data, "sa1 interrupt data")
			assert.True(t, info.WasInterrupted)
			assert.Nil(t, info.InterruptState)
			assert.True(t, info.IsResumeTarget)
			assert.Equal(t, "resume sa1", info.ResumeData)
			iter, generator := NewAsyncIteratorPair[*AgentEvent]()
			generator.Close()
			return iter
		},
	} // interrupt once
	sa2 := &myAgent{
		name: "sa2",
		runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
			iter, generator := NewAsyncIteratorPair[*AgentEvent]()

			intEvent := StatefulInterrupt(ctx, "sa2 interrupt data", "sa2 interrupt")
			intEvent.Action.Interrupted.Data = "sa2 interrupt data"
			generator.Send(intEvent)
			generator.Close()
			return iter
		},
		resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {
			assert.Equal(t, info.InterruptInfo.Data, "sa2 interrupt data")
			assert.True(t, info.WasInterrupted)
			assert.NotNil(t, info.InterruptState)
			assert.Equal(t, "sa2 interrupt", info.InterruptState)

			assert.True(t, info.IsResumeTarget)
			assert.NotNil(t, info.ResumeData)
			assert.Equal(t, "resume sa2", info.ResumeData)
			iter, generator := NewAsyncIteratorPair[*AgentEvent]()
			generator.Close()
			return iter
		},
	} // interrupt once
	sa3 := &myAgent{
		name: "sa3",
		runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
			iter, generator := NewAsyncIteratorPair[*AgentEvent]()
			generator.Send(&AgentEvent{
				AgentName: "sa3",
				Output: &AgentOutput{
					MessageOutput: &MessageVariant{
						Message: schema.UserMessage("sa3 completed"),
					},
				},
			})
			generator.Close()
			return iter
		},
	} // won't interrupt
	sa4 := &myAgent{
		name: "sa4",
		runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
			iter, generator := NewAsyncIteratorPair[*AgentEvent]()
			generator.Send(&AgentEvent{
				AgentName: "sa4",
				Output: &AgentOutput{
					MessageOutput: &MessageVariant{
						Message: schema.UserMessage("sa4 completed"),
					},
				},
			})
			generator.Close()
			return iter
		},
	} // won't interrupt

	firstInterruptEvent := &AgentEvent{
		AgentName: "sa1",
		RunPath:   []RunStep{{"sequential"}, {"sa1"}},
		Action: &AgentAction{
			Interrupted: &InterruptInfo{
				Data: &WorkflowInterruptInfo{
					OrigInput: &AgentInput{
						Messages: []Message{schema.UserMessage("hello world")},
					},
					SequentialInterruptIndex: 0,
					SequentialInterruptInfo: &InterruptInfo{
						Data: "sa1 interrupt data",
					},
					LoopIterations: 0,
				},
				InterruptContexts: []*InterruptCtx{
					{
						ID:   "agent:sequential;agent:sa1",
						Info: "sa1 interrupt data",
						Address: Address{
							{
								ID:   "sequential",
								Type: AddressSegmentAgent,
							},
							{
								ID:   "sa1",
								Type: AddressSegmentAgent,
							},
						},
						IsRootCause: true,
						Parent: &InterruptCtx{
							ID:   "agent:sequential",
							Info: "Sequential workflow interrupted",
							Address: Address{
								{
									ID:   "sequential",
									Type: AddressSegmentAgent,
								},
							},
						},
					},
				},
			},
		},
	}
	_ = firstInterruptEvent
	secondInterruptEvent := &AgentEvent{
		AgentName: "sa2",
		RunPath:   []RunStep{{"sequential"}, {"sa1"}, {"sa2"}},
		Action: &AgentAction{
			Interrupted: &InterruptInfo{
				Data: &WorkflowInterruptInfo{
					OrigInput: &AgentInput{
						Messages: []Message{schema.UserMessage("hello world")},
					},
					SequentialInterruptIndex: 1,
					SequentialInterruptInfo: &InterruptInfo{
						Data: "sa2 interrupt data",
					},
				},
				InterruptContexts: []*InterruptCtx{
					{
						ID:   "agent:sequential;agent:sa1;agent:sa2",
						Info: "sa2 interrupt data",
						Address: Address{
							{
								ID:   "sequential",
								Type: AddressSegmentAgent,
							},
							{
								ID:   "sa2",
								Type: AddressSegmentAgent,
							},
						},
						IsRootCause: true,
						Parent: &InterruptCtx{
							ID:   "agent:sequential",
							Info: "Sequential workflow interrupted",
							Address: Address{
								{
									ID:   "sequential",
									Type: AddressSegmentAgent,
								},
							},
						},
					},
				},
			},
		},
	}
	_ = secondInterruptEvent
	messageEvents := []*AgentEvent{
		{
			AgentName: "sa3",
			RunPath:   []RunStep{{"sequential"}, {"sa1"}, {"sa2"}, {"sa3"}},
			Output: &AgentOutput{
				MessageOutput: &MessageVariant{
					Message: schema.UserMessage("sa3 completed"),
				},
			},
		},
		{
			AgentName: "sa4",
			RunPath:   []RunStep{{"sequential"}, {"sa1"}, {"sa2"}, {"sa3"}, {"sa4"}},
			Output: &AgentOutput{
				MessageOutput: &MessageVariant{
					Message: schema.UserMessage("sa4 completed"),
				},
			},
		},
	}
	_ = messageEvents

	t.Run("test sequential workflow agent", func(t *testing.T) {

		// sequential
		a, err := NewSequentialAgent(ctx, &SequentialAgentConfig{
			Name:        "sequential",
			Description: "sequential agent",
			SubAgents:   []Agent{sa1, sa2, sa3, sa4},
		})
		assert.NoError(t, err)
		runner := NewRunner(ctx, RunnerConfig{
			Agent:           a,
			CheckPointStore: newMyStore(),
		})
		var events []*AgentEvent
		iter := runner.Query(ctx, "hello world", WithCheckPointID("sequential-1"))
		for {
			event, ok := iter.Next()
			if !ok {
				break
			}
			events = append(events, event)
		}

		assert.Equal(t, 1, len(events))
		assert.Equal(t, firstInterruptEvent.AgentName, events[0].AgentName)
		assert.Equal(t, firstInterruptEvent.RunPath, events[0].RunPath)
		assert.True(t, events[0].Action.Interrupted.InterruptContexts[0].EqualsWithoutID(firstInterruptEvent.Action.Interrupted.InterruptContexts[0]))
		interruptID1 := events[0].Action.Interrupted.InterruptContexts[0].ID
		events = []*AgentEvent{}

		// Resume after sa1 interrupt
		iter, err = runner.ResumeWithParams(ctx, "sequential-1", &ResumeParams{
			Targets: map[string]any{
				interruptID1: "resume sa1",
			},
		})
		assert.NoError(t, err)
		for {
			event, ok := iter.Next()
			if !ok {
				break
			}
			events = append(events, event)
		}

		assert.Equal(t, 1, len(events))
		assert.Equal(t, secondInterruptEvent.AgentName, events[0].AgentName)
		assert.Equal(t, secondInterruptEvent.RunPath, events[0].RunPath)
		assert.True(t, events[0].Action.Interrupted.InterruptContexts[0].
			EqualsWithoutID(secondInterruptEvent.Action.Interrupted.InterruptContexts[0]))
		interruptID2 := events[0].Action.Interrupted.InterruptContexts[0].ID
		events = []*AgentEvent{}

		// Resume after sa2 interrupt
		iter, err = runner.ResumeWithParams(ctx, "sequential-1", &ResumeParams{
			Targets: map[string]any{
				interruptID2: "resume sa2",
			},
		})
		assert.NoError(t, err)
		for {
			event, ok := iter.Next()
			if !ok {
				break
			}
			events = append(events, event)
		}

		assert.Equal(t, 2, len(events))
		assert.Equal(t, messageEvents, events)
	})

	t.Run("test loop workflow agent", func(t *testing.T) {
		// loop
		a, err := NewLoopAgent(ctx, &LoopAgentConfig{
			Name:          "loop",
			SubAgents:     []Agent{sa1, sa2, sa3, sa4},
			MaxIterations: 2,
		})
		assert.NoError(t, err)
		runner := NewRunner(ctx, RunnerConfig{
			Agent:           a,
			CheckPointStore: newMyStore(),
		})
		var events []*AgentEvent
		iter := runner.Query(ctx, "hello world", WithCheckPointID("loop-1"))
		for {
			event, ok := iter.Next()
			if !ok {
				break
			}
			events = append(events, event)
		}

		loopFirstInterruptEvent := &AgentEvent{
			AgentName: "sa1",
			RunPath:   []RunStep{{"loop"}, {"sa1"}},
			Action: &AgentAction{
				Interrupted: &InterruptInfo{
					Data: &WorkflowInterruptInfo{
						OrigInput: &AgentInput{
							Messages: []Message{schema.UserMessage("hello world")},
						},
						SequentialInterruptIndex: 0,
						SequentialInterruptInfo: &InterruptInfo{
							Data: "sa1 interrupt data",
						},
						LoopIterations: 0,
					},
					InterruptContexts: []*InterruptCtx{
						{
							ID:   "agent:loop;agent:sa1",
							Info: "sa1 interrupt data",
							Address: Address{
								{
									ID:   "loop",
									Type: AddressSegmentAgent,
								},
								{
									ID:   "sa1",
									Type: AddressSegmentAgent,
								},
							},
							IsRootCause: true,
							Parent: &InterruptCtx{
								ID:   "agent:loop",
								Info: "Loop workflow interrupted",
								Address: Address{
									{
										ID:   "loop",
										Type: AddressSegmentAgent,
									},
								},
							},
						},
					},
				},
			},
		}
		assert.Equal(t, 1, len(events))
		assert.Equal(t, loopFirstInterruptEvent.AgentName, events[0].AgentName)
		assert.Equal(t, loopFirstInterruptEvent.RunPath, events[0].RunPath)
		assert.True(t, events[0].Action.Interrupted.InterruptContexts[0].EqualsWithoutID(loopFirstInterruptEvent.Action.Interrupted.InterruptContexts[0]))
		loopInterruptID1 := events[0].Action.Interrupted.InterruptContexts[0].ID
		events = []*AgentEvent{}

		// Resume after sa1 interrupt
		iter, err = runner.ResumeWithParams(ctx, "loop-1", &ResumeParams{
			Targets: map[string]any{
				loopInterruptID1: "resume sa1",
			},
		})
		assert.NoError(t, err)
		for {
			event, ok := iter.Next()
			if !ok {
				break
			}
			events = append(events, event)
		}

		loopSecondInterruptEvent := &AgentEvent{
			AgentName: "sa2",
			RunPath:   []RunStep{{"loop"}, {"sa1"}, {"sa2"}},
			Action: &AgentAction{
				Interrupted: &InterruptInfo{
					Data: &WorkflowInterruptInfo{
						OrigInput: &AgentInput{
							Messages: []Message{schema.UserMessage("hello world")},
						},
						SequentialInterruptIndex: 1,
						SequentialInterruptInfo: &InterruptInfo{
							Data: "sa2 interrupt data",
						},
						LoopIterations: 0,
					},
					InterruptContexts: []*InterruptCtx{
						{
							ID:   "agent:loop;agent:sa1;agent:sa2",
							Info: "sa2 interrupt data",
							Address: Address{
								{
									ID:   "loop",
									Type: AddressSegmentAgent,
								},
								{
									ID:   "sa2",
									Type: AddressSegmentAgent,
								},
							},
							IsRootCause: true,
							Parent: &InterruptCtx{
								ID:   "agent:loop",
								Info: "Loop workflow interrupted",
								Address: Address{
									{
										ID:   "loop",
										Type: AddressSegmentAgent,
									},
								},
							},
						},
					},
				},
			},
		}
		assert.Equal(t, 1, len(events))
		assert.Equal(t, loopSecondInterruptEvent.AgentName, events[0].AgentName)
		assert.Equal(t, loopSecondInterruptEvent.RunPath, events[0].RunPath)
		assert.True(t, events[0].Action.Interrupted.InterruptContexts[0].EqualsWithoutID(loopSecondInterruptEvent.Action.Interrupted.InterruptContexts[0]))
		loopInterruptID2 := events[0].Action.Interrupted.InterruptContexts[0].ID
		events = []*AgentEvent{}

		// Resume after sa2 interrupt
		iter, err = runner.ResumeWithParams(ctx, "loop-1", &ResumeParams{
			Targets: map[string]any{
				loopInterruptID2: "resume sa2",
			},
		})
		assert.NoError(t, err)
		for {
			event, ok := iter.Next()
			if !ok {
				break
			}
			events = append(events, event)
		}

		loopThirdInterruptEvent := &AgentEvent{
			AgentName: "sa1",
			RunPath:   []RunStep{{"loop"}, {"sa1"}, {"sa2"}, {"sa3"}, {"sa4"}, {"sa1"}},
			Action: &AgentAction{
				Interrupted: &InterruptInfo{
					Data: &WorkflowInterruptInfo{
						OrigInput: &AgentInput{
							Messages: []Message{schema.UserMessage("hello world")},
						},
						SequentialInterruptIndex: 0,
						SequentialInterruptInfo: &InterruptInfo{
							Data: "sa1 interrupt data",
						},
						LoopIterations: 1,
					},
					InterruptContexts: []*InterruptCtx{
						{
							ID:   "agent:loop;agent:sa1;agent:sa2;agent:sa3;agent:sa4;agent:sa1",
							Info: "sa1 interrupt data",
							Address: Address{
								{
									ID:   "loop",
									Type: AddressSegmentAgent,
								},
								{
									ID:   "sa1",
									Type: AddressSegmentAgent,
								},
							},
							IsRootCause: true,
							Parent: &InterruptCtx{
								ID:   "agent:loop",
								Info: "Loop workflow interrupted",
								Address: Address{
									{
										ID:   "loop",
										Type: AddressSegmentAgent,
									},
								},
							},
						},
					},
				},
			},
		}

		loopFourthInterruptEvent := &AgentEvent{
			AgentName: "sa2",
			RunPath:   []RunStep{{"loop"}, {"sa1"}, {"sa2"}, {"sa3"}, {"sa4"}, {"sa1"}, {"sa2"}},
			Action: &AgentAction{
				Interrupted: &InterruptInfo{
					Data: &WorkflowInterruptInfo{
						OrigInput: &AgentInput{
							Messages: []Message{schema.UserMessage("hello world")},
						},
						SequentialInterruptIndex: 1,
						SequentialInterruptInfo: &InterruptInfo{
							Data: "sa2 interrupt data",
						},
						LoopIterations: 1,
					},
					InterruptContexts: []*InterruptCtx{
						{
							ID:   "agent:loop;agent:sa1;agent:sa2;agent:sa3;agent:sa4;agent:sa1;agent:sa2",
							Info: "sa2 interrupt data",
							Address: Address{
								{
									ID:   "loop",
									Type: AddressSegmentAgent,
								},
								{
									ID:   "sa2",
									Type: AddressSegmentAgent,
								},
							},
							IsRootCause: true,
							Parent: &InterruptCtx{
								ID:   "agent:loop",
								Info: "Loop workflow interrupted",
								Address: Address{
									{
										ID:   "loop",
										Type: AddressSegmentAgent,
									},
								},
							},
						},
					},
				},
			},
		}

		loopMessageEvents := []*AgentEvent{
			{
				AgentName: "sa3",
				RunPath:   []RunStep{{"loop"}, {"sa1"}, {"sa2"}, {"sa3"}},
				Output: &AgentOutput{
					MessageOutput: &MessageVariant{
						Message: schema.UserMessage("sa3 completed"),
					},
				},
			},
			{
				AgentName: "sa4",
				RunPath:   []RunStep{{"loop"}, {"sa1"}, {"sa2"}, {"sa3"}, {"sa4"}},
				Output: &AgentOutput{
					MessageOutput: &MessageVariant{
						Message: schema.UserMessage("sa4 completed"),
					},
				},
			},
			loopThirdInterruptEvent,
		}
		assert.Equal(t, 3, len(events))
		// Check the first two message events
		assert.Equal(t, loopMessageEvents[0].AgentName, events[0].AgentName)
		assert.Equal(t, loopMessageEvents[0].RunPath, events[0].RunPath)
		assert.Equal(t, loopMessageEvents[0].Output.MessageOutput.Message.Content, events[0].Output.MessageOutput.Message.Content)

		assert.Equal(t, loopMessageEvents[1].AgentName, events[1].AgentName)
		assert.Equal(t, loopMessageEvents[1].RunPath, events[1].RunPath)
		assert.Equal(t, loopMessageEvents[1].Output.MessageOutput.Message.Content, events[1].Output.MessageOutput.Message.Content)

		// Check the third interrupt event using EqualsWithoutID
		assert.Equal(t, loopMessageEvents[2].AgentName, events[2].AgentName)
		assert.Equal(t, loopMessageEvents[2].RunPath, events[2].RunPath)
		assert.True(t, events[2].Action.Interrupted.InterruptContexts[0].EqualsWithoutID(loopMessageEvents[2].Action.Interrupted.InterruptContexts[0]))
		loopInterruptID3 := events[2].Action.Interrupted.InterruptContexts[0].ID
		events = []*AgentEvent{}

		// Resume after third interrupt
		iter, err = runner.ResumeWithParams(ctx, "loop-1", &ResumeParams{
			Targets: map[string]any{
				loopInterruptID3: "resume sa1",
			},
		})
		assert.NoError(t, err)
		for {
			event, ok := iter.Next()
			if !ok {
				break
			}
			events = append(events, event)
		}
		assert.Equal(t, 1, len(events))
		assert.Equal(t, loopFourthInterruptEvent.AgentName, events[0].AgentName)
		assert.Equal(t, loopFourthInterruptEvent.RunPath, events[0].RunPath)
		assert.True(t, events[0].Action.Interrupted.InterruptContexts[0].EqualsWithoutID(loopFourthInterruptEvent.Action.Interrupted.InterruptContexts[0]))
		loopInterruptID4 := events[0].Action.Interrupted.InterruptContexts[0].ID
		events = []*AgentEvent{}

		// Resume after fourth interrupt
		iter, err = runner.ResumeWithParams(ctx, "loop-1", &ResumeParams{
			Targets: map[string]any{
				loopInterruptID4: "resume sa2",
			},
		})
		assert.NoError(t, err)
		for {
			event, ok := iter.Next()
			if !ok {
				break
			}
			events = append(events, event)
		}
		loopFinalMessageEvents := []*AgentEvent{
			{
				AgentName: "sa3",
				RunPath:   []RunStep{{"loop"}, {"sa1"}, {"sa2"}, {"sa3"}, {"sa4"}, {"sa1"}, {"sa2"}, {"sa3"}},
				Output: &AgentOutput{
					MessageOutput: &MessageVariant{
						Message: schema.UserMessage("sa3 completed"),
					},
				},
			},
			{
				AgentName: "sa4",
				RunPath:   []RunStep{{"loop"}, {"sa1"}, {"sa2"}, {"sa3"}, {"sa4"}, {"sa1"}, {"sa2"}, {"sa3"}, {"sa4"}},
				Output: &AgentOutput{
					MessageOutput: &MessageVariant{
						Message: schema.UserMessage("sa4 completed"),
					},
				},
			},
		}
		assert.Equal(t, 2, len(events))
		assert.Equal(t, loopFinalMessageEvents, events)
	})

	t.Run("test parallel workflow agent", func(t *testing.T) {
		// parallel
		a, err := NewParallelAgent(ctx, &ParallelAgentConfig{
			Name:      "parallel agent",
			SubAgents: []Agent{sa1, sa2, sa3, sa4},
		})
		assert.NoError(t, err)
		runner := NewRunner(ctx, RunnerConfig{
			Agent:           a,
			CheckPointStore: newMyStore(),
		})
		iter := runner.Query(ctx, "hello world", WithCheckPointID("1"))
		var (
			events         []*AgentEvent
			interruptEvent *AgentEvent
		)

		for {
			event, ok := iter.Next()
			if !ok {
				break
			}
			if event.Action != nil && event.Action.Interrupted != nil {
				interruptEvent = event
				continue
			}
			events = append(events, event)
		}
		assert.Equal(t, 2, len(events))

		// Debug: Print actual events to see what we're getting
		for i, event := range events {
			t.Logf("Event %d: AgentName=%s, RunPath=%v, Output=%v", i, event.AgentName, event.RunPath, event.Output)
		}

		// Define parallel message events separately
		parallelMessageEvents := []*AgentEvent{
			{
				AgentName: "sa4",
				RunPath:   []RunStep{{"parallel agent"}, {"sa4"}},
				Output: &AgentOutput{
					MessageOutput: &MessageVariant{
						Message: schema.UserMessage("sa4 completed"),
					},
				},
			},
			{
				AgentName: "sa3",
				RunPath:   []RunStep{{"parallel agent"}, {"sa3"}},
				Output: &AgentOutput{
					MessageOutput: &MessageVariant{
						Message: schema.UserMessage("sa3 completed"),
					},
				},
			},
		}

		assert.Contains(t, events, parallelMessageEvents[0])
		assert.Contains(t, events, parallelMessageEvents[1])

		assert.NotNil(t, interruptEvent)
		assert.Equal(t, "parallel agent", interruptEvent.AgentName)
		assert.Equal(t, []RunStep{{"parallel agent"}}, interruptEvent.RunPath)
		assert.NotNil(t, interruptEvent.Action.Interrupted)
		wii, ok := interruptEvent.Action.Interrupted.Data.(*WorkflowInterruptInfo)
		assert.True(t, ok)
		assert.Equal(t, 2, len(wii.ParallelInterruptInfo))

		var sa1Found, sa2Found bool
		for _, info := range wii.ParallelInterruptInfo {
			switch info.Data {
			case "sa1 interrupt data":
				sa1Found = true
			case "sa2 interrupt data":
				sa2Found = true
			}
		}
		assert.True(t, sa1Found)
		assert.True(t, sa2Found)

		var sa1InfoFound, sa2InfoFound bool
		for _, ctx := range interruptEvent.Action.Interrupted.InterruptContexts {
			switch ctx.Info {
			case "sa1 interrupt data":
				sa1InfoFound = true
			case "sa2 interrupt data":
				sa2InfoFound = true
			}
		}

		assert.Equal(t, 2, len(interruptEvent.Action.Interrupted.InterruptContexts))
		assert.True(t, sa1InfoFound)
		assert.True(t, sa2InfoFound)

		var parallelInterruptID1, parallelInterruptID2 string
		for _, ctx := range interruptEvent.Action.Interrupted.InterruptContexts {
			switch ctx.Info {
			case "sa1 interrupt data":
				parallelInterruptID1 = ctx.ID
			case "sa2 interrupt data":
				parallelInterruptID2 = ctx.ID
			}
		}
		assert.NotEmpty(t, parallelInterruptID1)
		assert.NotEmpty(t, parallelInterruptID2)

		iter, err = runner.ResumeWithParams(ctx, "1", &ResumeParams{
			Targets: map[string]any{
				parallelInterruptID1: "resume sa1",
				parallelInterruptID2: "resume sa2",
			},
		})
		assert.NoError(t, err)
		_, ok = iter.Next()
		assert.False(t, ok)
	})
}

func TestChatModelInterrupt(t *testing.T) {
	ctx := context.Background()
	a, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
		Name:        "name",
		Description: "description",
		Instruction: "instruction",
		Model: &myModel{
			validator: func(i int, messages []*schema.Message) bool {
				if i > 0 && (len(messages) != 4 || messages[2].Content != "new user message") {
					return false
				}
				return true
			},
			messages: []*schema.Message{
				schema.AssistantMessage("", []schema.ToolCall{
					{
						ID: "1",
						Function: schema.FunctionCall{
							Name:      "tool1",
							Arguments: "arguments",
						},
					},
				}),
				schema.AssistantMessage("completed", nil),
			},
		},
		ToolsConfig: ToolsConfig{
			ToolsNodeConfig: compose.ToolsNodeConfig{
				Tools: []tool.BaseTool{&myTool1{}},
			},
		},
	})
	assert.NoError(t, err)
	runner := NewRunner(ctx, RunnerConfig{
		Agent:           a,
		CheckPointStore: newMyStore(),
	})
	iter := runner.Query(ctx, "hello world", WithCheckPointID("1"))
	_, ok := iter.Next()
	assert.True(t, ok)
	event, ok := iter.Next()
	assert.True(t, ok)
	assert.NoError(t, event.Err)
	assert.NotNil(t, event.Action.Interrupted)
	assert.Equal(t, 1, len(event.Action.Interrupted.InterruptContexts))
	assert.Equal(t, Address{
		{Type: AddressSegmentAgent, ID: "name"},
		{Type: AddressSegmentTool, ID: "tool1", SubID: "1"},
	}, event.Action.Interrupted.InterruptContexts[0].Address)

	var (
		chatModelAgentID string
		toolID           string
	)

	intCtx := event.Action.Interrupted.InterruptContexts[0]
	for intCtx != nil {
		if intCtx.Address[len(intCtx.Address)-1].Type == AddressSegmentTool {
			toolID = intCtx.ID
		} else if intCtx.Address[len(intCtx.Address)-1].Type == AddressSegmentAgent {
			chatModelAgentID = intCtx.ID
		}
		intCtx = intCtx.Parent
	}

	_, ok = iter.Next()
	assert.False(t, ok)

	iter, err = runner.ResumeWithParams(ctx, "1", &ResumeParams{
		Targets: map[string]any{
			chatModelAgentID: &ChatModelAgentResumeData{
				HistoryModifier: func(ctx context.Context, history []Message) []Message {
					history[2].Content = "new user message"
					return history
				},
			},
			toolID: "tool resume result",
		},
	})
	assert.NoError(t, err)
	event, ok = iter.Next()
	assert.True(t, ok)
	assert.NoError(t, event.Err)
	assert.Equal(t, event.Output.MessageOutput.Message.Content, "tool resume result")
	event, ok = iter.Next()
	assert.True(t, ok)
	assert.NoError(t, event.Err)
	assert.Equal(t, event.Output.MessageOutput.Message.Content, "completed")
}

func TestChatModelAgentToolInterrupt(t *testing.T) {
	sa := &myAgent{
		runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
			iter, generator := NewAsyncIteratorPair[*AgentEvent]()
			intAct := Interrupt(ctx, "hello world")
			intAct.Action.Interrupted.Data = "hello world"
			generator.Send(intAct)
			generator.Close()
			return iter
		},
		resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {
			assert.NotNil(t, info)
			assert.False(t, info.EnableStreaming)

			if !info.IsResumeTarget {
				iter, generator := NewAsyncIteratorPair[*AgentEvent]()
				intAct := Interrupt(ctx, "interrupt again")
				intAct.Action.Interrupted.Data = "interrupt again"
				generator.Send(intAct)
				generator.Close()
				return iter
			}

			assert.NotNil(t, info.ResumeData)
			assert.Equal(t, "resume sa", info.ResumeData)

			iter, generator := NewAsyncIteratorPair[*AgentEvent]()
			generator.Send(&AgentEvent{Output: &AgentOutput{MessageOutput: &MessageVariant{Message: schema.UserMessage(fmt.Sprintf("my agent completed with data %s", info.ResumeData))}}})
			generator.Close()
			return iter
		},
	}
	ctx := context.Background()
	a, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
		Name:        "name",
		Description: "description",
		Instruction: "instruction",
		Model: &myModel{
			messages: []*schema.Message{
				schema.AssistantMessage("", []schema.ToolCall{
					{
						ID: "1",
						Function: schema.FunctionCall{
							Name:      "myAgent",
							Arguments: "{\"request\":\"123\"}",
						},
					},
				}),
				schema.AssistantMessage("completed", nil),
			},
		},
		ToolsConfig: ToolsConfig{
			ToolsNodeConfig: compose.ToolsNodeConfig{
				Tools: []tool.BaseTool{NewAgentTool(ctx, sa)},
			},
		},
	})
	assert.NoError(t, err)
	runner := NewRunner(ctx, RunnerConfig{
		Agent:           a,
		CheckPointStore: newMyStore(),
	})

	iter := runner.Query(ctx, "hello world", WithCheckPointID("1"))
	_, ok := iter.Next()
	assert.True(t, ok)
	event, ok := iter.Next()
	assert.True(t, ok)
	assert.NoError(t, event.Err)
	assert.NotNil(t, event.Action.Interrupted)
	_, ok = iter.Next()
	assert.False(t, ok)

	iter, err = runner.Resume(ctx, "1")
	assert.NoError(t, err)
	event, ok = iter.Next()
	assert.True(t, ok)
	assert.NoError(t, event.Err)
	assert.NotNil(t, event.Action.Interrupted)
	assert.Equal(t, 1, len(event.Action.Interrupted.InterruptContexts))
	for _, ctx := range event.Action.Interrupted.InterruptContexts {
		if ctx.IsRootCause {
			assert.Equal(t, Address{
				{Type: AddressSegmentAgent, ID: "name"},
				{Type: AddressSegmentTool, ID: "myAgent", SubID: "1"},
				{Type: AddressSegmentAgent, ID: "myAgent"},
			}, ctx.Address)
			assert.Equal(t, "interrupt again", ctx.Info)
		}
	}

	var toolInterruptID string
	for _, ctx := range event.Action.Interrupted.InterruptContexts {
		if ctx.IsRootCause {
			toolInterruptID = ctx.ID
			break
		}
	}
	assert.NotEmpty(t, toolInterruptID)

	_, ok = iter.Next()
	assert.False(t, ok)

	iter, err = runner.ResumeWithParams(ctx, "1", &ResumeParams{
		Targets: map[string]any{
			toolInterruptID: "resume sa",
		},
	})
	assert.NoError(t, err)
	event, ok = iter.Next()
	assert.True(t, ok)
	assert.NoError(t, event.Err)
	assert.Equal(t, event.Output.MessageOutput.Message.Content, "my agent completed with data resume sa")
	event, ok = iter.Next()
	assert.True(t, ok)
	assert.NoError(t, event.Err)
	assert.Equal(t, event.Output.MessageOutput.Message.Content, "completed")
	_, ok = iter.Next()
	assert.False(t, ok)
}

func newMyStore() *myStore {
	return &myStore{
		m: map[string][]byte{},
	}
}

type myStore struct {
	m map[string][]byte
}

func (m *myStore) Set(_ context.Context, key string, value []byte) error {
	m.m[key] = value
	return nil
}

func (m *myStore) Get(_ context.Context, key string) ([]byte, bool, error) {
	v, ok := m.m[key]
	return v, ok, nil
}

type myAgentOptions struct {
	value string
}

func withValue(value string) AgentRunOption {
	return WrapImplSpecificOptFn(func(t *myAgentOptions) {
		t.value = value
	})
}

type myAgent struct {
	name     string
	runFn    func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent]
	resumeFn func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent]
}

func (m *myAgent) Name(_ context.Context) string {
	if len(m.name) > 0 {
		return m.name
	}
	return "myAgent"
}

func (m *myAgent) Description(_ context.Context) string {
	return "myAgent description"
}

func (m *myAgent) Run(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
	return m.runFn(ctx, input, options...)
}

func (m *myAgent) Resume(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {
	return m.resumeFn(ctx, info, opts...)
}

type myModel struct {
	times     int
	messages  []*schema.Message
	validator func(int, []*schema.Message) bool
}

func (m *myModel) Generate(_ context.Context, input []*schema.Message, _ ...model.Option) (*schema.Message, error) {
	if m.validator != nil && !m.validator(m.times, input) {
		return nil, errors.New("invalid input")
	}
	if m.times >= len(m.messages) {
		return nil, errors.New("exceeded max number of messages")
	}
	t := m.times
	m.times++
	return m.messages[t], nil
}

func (m *myModel) Stream(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
	panic("implement me")
}

func (m *myModel) WithTools(_ []*schema.ToolInfo) (model.ToolCallingChatModel, error) {
	return m, nil
}

type myTool1 struct{}

func (m *myTool1) Info(_ context.Context) (*schema.ToolInfo, error) {
	return &schema.ToolInfo{
		Name: "tool1",
		Desc: "desc",
	}, nil
}

func (m *myTool1) InvokableRun(ctx context.Context, _ string, _ ...tool.Option) (string, error) {
	if wasInterrupted, _, _ := tool.GetInterruptState[any](ctx); !wasInterrupted {
		return "", tool.Interrupt(ctx, nil)
	}

	if isResumeFlow, hasResumeData, data := tool.GetResumeContext[string](ctx); !isResumeFlow {
		return "", tool.Interrupt(ctx, nil)
	} else if hasResumeData {
		return data, nil
	}

	return "result", nil
}

func TestCyclicalAgentInterrupt(t *testing.T) {
	ctx := context.Background()

	var agentA, agentB, agentC Agent

	// agentC interrupts
	agentC = &myAgent{
		name: "C",
		runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
			iter, generator := NewAsyncIteratorPair[*AgentEvent]()
			intAct := Interrupt(ctx, "interrupt from C")
			generator.Send(intAct)
			generator.Close()
			return iter
		},
		resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {
			assert.True(t, info.IsResumeTarget)
			assert.NotNil(t, info.ResumeData)
			assert.Equal(t, "resume C", info.ResumeData)
			iter, generator := NewAsyncIteratorPair[*AgentEvent]()
			generator.Send(&AgentEvent{
				AgentName: "C",
				Output: &AgentOutput{
					MessageOutput: &MessageVariant{Message: schema.UserMessage("C completed")},
				},
			})
			generator.Close()
			return iter
		},
	}

	// agentB transfers back to its parent A
	agentB = &myAgent{
		name: "B",
		runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
			iter, generator := NewAsyncIteratorPair[*AgentEvent]()
			generator.Send(&AgentEvent{
				AgentName: "B",
				Action: &AgentAction{
					TransferToAgent: &TransferToAgentAction{
						DestAgentName: "A", // Transfer back to parent
					},
				},
			})
			generator.Close()
			return iter
		},
	}

	// agentA is the parent, orchestrating the A->B->A->C flow
	agentA = &myAgent{
		name: "A",
		runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
			runCtx := getRunCtx(ctx)
			iter, generator := NewAsyncIteratorPair[*AgentEvent]()

			// If the last agent was B, we are in the A->B->A path, so transfer to C.
			// Otherwise, it's the first run, transfer to B.
			dest := "B"
			if len(runCtx.RunPath) > 1 && runCtx.RunPath[len(runCtx.RunPath)-2].agentName == "B" {
				dest = "C"
			}

			generator.Send(&AgentEvent{
				AgentName: "A",
				Action: &AgentAction{
					TransferToAgent: &TransferToAgentAction{
						DestAgentName: dest,
					},
				},
			})
			generator.Close()
			return iter
		},
	}

	// Set up the hierarchy: A is parent of B and C.
	agentA, err := SetSubAgents(ctx, agentA, []Agent{agentB, agentC})
	assert.NoError(t, err)

	// Run the test
	runner := NewRunner(ctx, RunnerConfig{
		Agent:           agentA,
		CheckPointStore: newMyStore(),
	})
	iter := runner.Query(ctx, "start", WithCheckPointID("cyclical-1"))

	var events []*AgentEvent
	for {
		event, ok := iter.Next()
		if !ok {
			break
		}
		events = append(events, event)
	}

	// We expect 3 transfer events (A->B, B->A, A->C) and 1 interrupt event from C.
	assert.Equal(t, 4, len(events))

	interruptEvent := events[3]
	assert.NotNil(t, interruptEvent.Action.Interrupted)
	assert.Equal(t, "C", interruptEvent.AgentName)

	// Check the interrupt context
	assert.Equal(t, 1, len(interruptEvent.Action.Interrupted.InterruptContexts))
	interruptCtx := interruptEvent.Action.Interrupted.InterruptContexts[0]
	assert.True(t, interruptCtx.IsRootCause)
	assert.Equal(t, "interrupt from C", interruptCtx.Info)

	expectedAddr := Address{
		{Type: AddressSegmentAgent, ID: "A"},
		{Type: AddressSegmentAgent, ID: "B"},
		{Type: AddressSegmentAgent, ID: "A"},
		{Type: AddressSegmentAgent, ID: "C"},
	}
	assert.Equal(t, expectedAddr, interruptCtx.Address)
	assert.NotEmpty(t, interruptCtx.ID)

	// Resume the execution
	iter, err = runner.ResumeWithParams(ctx, "cyclical-1", &ResumeParams{
		Targets: map[string]any{
			interruptCtx.ID: "resume C",
		},
	})
	assert.NoError(t, err)

	events = []*AgentEvent{}
	for {
		event, ok := iter.Next()
		if !ok {
			break
		}
		events = append(events, event)
	}

	// We expect one output event from C
	assert.Equal(t, 1, len(events))
	assert.Equal(t, "C completed", events[0].Output.MessageOutput.Message.Content)
}

// myStatefulTool is a tool that can interrupt and has internal state to track invocations.

type myStatefulTool struct {
	name string
	t    *testing.T
}

func (m *myStatefulTool) Info(_ context.Context) (*schema.ToolInfo, error) {
	return &schema.ToolInfo{
		Name: m.name,
		Desc: "desc",
	}, nil
}

type myStatefulToolState struct {
	InterruptCount int
}

func init() {
	schema.Register[myStatefulToolState]()
}

func (m *myStatefulTool) InvokableRun(ctx context.Context, _ string, _ ...tool.Option) (string, error) {
	wasInterrupted, hasState, state := tool.GetInterruptState[myStatefulToolState](ctx)
	if !wasInterrupted {
		return "", tool.StatefulInterrupt(ctx, fmt.Sprintf("interrupt from %s", m.name), myStatefulToolState{InterruptCount: 1})
	}

	isResumeFlow, hasResumeData, data := tool.GetResumeContext[string](ctx)
	if !isResumeFlow || !hasResumeData {
		assert.True(m.t, hasState, "tool %s should have interrupt state on resume", m.name)
		return "", tool.StatefulInterrupt(ctx, fmt.Sprintf("interrupt from %s", m.name), myStatefulToolState{InterruptCount: state.InterruptCount + 1})
	}

	return data, nil
}

func TestChatModelParallelToolInterruptAndResume(t *testing.T) {
	ctx := context.Background()

	toolA := &myStatefulTool{name: "toolA", t: t}
	toolB := &myStatefulTool{name: "toolB", t: t}

	chatModel, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
		Name:        "ParallelToolAgent",
		Description: "An agent that uses parallel tools",
		Model: &myModel{
			messages: []*schema.Message{
				// 1. First model response: call toolA and toolB in parallel
				schema.AssistantMessage("", []schema.ToolCall{
					{ID: "1", Function: schema.FunctionCall{Name: "toolA", Arguments: "{}"}},
					{ID: "2", Function: schema.FunctionCall{Name: "toolB", Arguments: "{}"}},
				}),
				// 2. Second model response (after tools are resumed): call them again to check state
				schema.AssistantMessage("", []schema.ToolCall{
					{ID: "3", Function: schema.FunctionCall{Name: "toolA", Arguments: "{}"}},
					{ID: "4", Function: schema.FunctionCall{Name: "toolB", Arguments: "{}"}},
				}),
				// 3. Final completion
				schema.AssistantMessage("all done", nil),
			},
		},
		ToolsConfig: ToolsConfig{
			ToolsNodeConfig: compose.ToolsNodeConfig{
				Tools: []tool.BaseTool{toolA, toolB},
			},
		},
	})
	assert.NoError(t, err)

	runner := NewRunner(ctx, RunnerConfig{
		Agent:           chatModel,
		CheckPointStore: newMyStore(),
	})

	// 1. Initial query -> parallel interrupt from toolA and toolB
	iter := runner.Query(ctx, "start", WithCheckPointID("parallel-tool-test-1"))
	normalEvents, interruptEvent := consumeUntilInterrupt(iter)

	assert.Equal(t, 1, len(normalEvents))
	assert.NotNil(t, interruptEvent)
	assert.Equal(t, 2, len(interruptEvent.Action.Interrupted.InterruptContexts),
		"should have 2 interrupts")

	var toolAInterruptID, toolBInterruptID string
	for _, info := range interruptEvent.Action.Interrupted.InterruptContexts {
		switch info.Info {
		case "interrupt from toolA":
			toolAInterruptID = info.ID
			assert.True(t, info.IsRootCause)
		case "interrupt from toolB":
			toolBInterruptID = info.ID
			assert.True(t, info.IsRootCause)
		}
	}
	assert.NotEmpty(t, toolAInterruptID)
	assert.NotEmpty(t, toolBInterruptID)

	// 2. Resume, targeting only toolA. toolB should re-interrupt.
	iter, err = runner.ResumeWithParams(ctx, "parallel-tool-test-1", &ResumeParams{
		Targets: map[string]any{
			toolAInterruptID: "toolA resumed",
		},
	})
	assert.NoError(t, err)
	_, interruptEvent = consumeUntilInterrupt(iter)

	assert.NotNil(t, interruptEvent, "expected a re-interrupt from toolB")
	assert.Equal(t, 1, len(interruptEvent.Action.Interrupted.InterruptContexts),
		"should have 1 remaining interrupts")

	var rootCause *InterruptCtx
	for _, info := range interruptEvent.Action.Interrupted.InterruptContexts {
		if info.IsRootCause {
			rootCause = info
			break
		}
	}

	if rootCause == nil {
		t.Fatal("expected a root cause interrupt from toolB")
	}
	assert.Equal(t, "interrupt from toolB", rootCause.Info)
	toolBReInterruptID := rootCause.ID

	// 3. Resume the re-interrupted toolB. The agent should then call the tools again.
	iter, err = runner.ResumeWithParams(ctx, "parallel-tool-test-1", &ResumeParams{
		Targets: map[string]any{
			toolBReInterruptID: "toolB resumed",
		},
	})
	assert.NoError(t, err)

	// 4. Consume all final events. The internal assertions in the tools will check the wasInterrupted flag.
	// We expect to see the results of the second tool calls, and then the final agent completion.
	finalEvents, interruptEvent := consumeUntilInterrupt(iter)
	assert.Equal(t, 2, len(finalEvents))
	assert.NotNil(t, interruptEvent)
}

// TestNestedChatModelAgentWithAgentTool verifies that the shouldFire method correctly prevents
// duplicate event firing in nested ChatModelAgent scenarios (ChatModelAgent -> AgentTool -> ChatModelAgent).
// This ensures that only the inner agent's cbHandler fires, not the outer agent's.
func TestNestedChatModelAgentWithAgentTool(t *testing.T) {
	ctx := context.Background()

	// Create an interruptible tool for the inner agent
	innerTool := &myStatefulTool{name: "innerTool", t: t}

	// Create the inner ChatModelAgent that will be wrapped by AgentTool
	innerAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
		Name:        "InnerAgent",
		Description: "Inner agent with interruptible tool",
		Model: &myModel{
			messages: []*schema.Message{
				schema.AssistantMessage("", []schema.ToolCall{
					{ID: "1", Function: schema.FunctionCall{Name: "innerTool", Arguments: "{}"}},
				}),
				schema.AssistantMessage("inner agent completed", nil),
			},
		},
		ToolsConfig: ToolsConfig{
			ToolsNodeConfig: compose.ToolsNodeConfig{
				Tools: []tool.BaseTool{innerTool},
			},
		},
	})
	assert.NoError(t, err)

	// Wrap the inner agent in an AgentTool
	agentTool := NewAgentTool(ctx, innerAgent)

	// Create the outer ChatModelAgent that uses the AgentTool
	outerAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
		Name:        "OuterAgent",
		Description: "Outer agent with AgentTool containing inner agent",
		Model: &myModel{
			messages: []*schema.Message{
				schema.AssistantMessage("", []schema.ToolCall{
					{ID: "1", Function: schema.FunctionCall{Name: "InnerAgent", Arguments: "{}"}},
				}),
				schema.AssistantMessage("outer agent completed", nil),
			},
		},
		ToolsConfig: ToolsConfig{
			ToolsNodeConfig: compose.ToolsNodeConfig{
				Tools: []tool.BaseTool{agentTool},
			},
		},
	})
	assert.NoError(t, err)

	runner := NewRunner(ctx, RunnerConfig{
		Agent:           outerAgent,
		CheckPointStore: newMyStore(),
	})

	// Run the query - this should trigger the nested agent structure
	iter := runner.Query(ctx, "start", WithCheckPointID("nested-agent-test-1"))

	// Collect all events to verify no duplicates
	var allEvents []*AgentEvent
	var interruptEvent *AgentEvent

	for {
		event, ok := iter.Next()
		if !ok {
			break
		}

		if event.Action != nil && event.Action.Interrupted != nil {
			assert.Nil(t, interruptEvent)
			interruptEvent = event
		}

		allEvents = append(allEvents, event)
	}

	if interruptEvent == nil {
		t.Fatal("expected an interrupt event")
	}

	// Verify we got exactly one interrupt event (not duplicated)
	assert.NotNil(t, interruptEvent, "should have an interrupt event")
	assert.Equal(t, 1, len(interruptEvent.Action.Interrupted.InterruptContexts),
		"should have exactly one interrupt context")

	// Verify the interrupt comes from the inner tool, not duplicated
	interruptCtx := interruptEvent.Action.Interrupted.InterruptContexts[0]
	assert.True(t, interruptCtx.IsRootCause, "interrupt should be root cause")
	assert.Equal(t, "interrupt from innerTool", interruptCtx.Info)

	// Verify the address path shows the correct nested structure
	expectedAddress := Address{
		{Type: AddressSegmentAgent, ID: "OuterAgent"},
		{Type: AddressSegmentTool, ID: "InnerAgent", SubID: "1"},
		{Type: AddressSegmentAgent, ID: "InnerAgent"},
		{Type: AddressSegmentTool, ID: "innerTool", SubID: "1"},
	}
	assert.Equal(t, expectedAddress, interruptCtx.Address,
		"interrupt address should show correct nested structure")

	// Verify no duplicate events by checking agent names in events
	var agentNames []string
	for _, event := range allEvents {
		if event.AgentName != "" {
			agentNames = append(agentNames, event.AgentName)
		}
	}

	// Should only have events from the outer agent (the inner agent's events should be handled
	// by the AgentTool and not duplicated by the outer agent's cbHandler)
	for _, name := range agentNames {
		assert.Equal(t, "OuterAgent", name,
			"all events should come from OuterAgent, not duplicated from InnerAgent")
	}

	// Now resume the interrupt
	interruptID := interruptCtx.ID
	iter, err = runner.ResumeWithParams(ctx, "nested-agent-test-1", &ResumeParams{
		Targets: map[string]any{
			interruptID: "resume inner tool",
		},
	})
	assert.NoError(t, err)

	// Collect final events after resume
	var finalEvents []*AgentEvent
	for {
		event, ok := iter.Next()
		if !ok {
			break
		}
		finalEvents = append(finalEvents, event)
	}

	// Verify completion events
	assert.Greater(t, len(finalEvents), 0, "should have completion events after resume")

	// Check that we get the expected completion messages
	var foundInnerCompletion, foundOuterCompletion bool
	for _, event := range finalEvents {
		if event.Output != nil && event.Output.MessageOutput != nil {
			if event.Output.MessageOutput.Message != nil {
				content := event.Output.MessageOutput.Message.Content
				switch content {
				case "inner agent completed":
					foundInnerCompletion = true
				case "outer agent completed":
					foundOuterCompletion = true
				}
			}
		}
	}

	assert.True(t, foundInnerCompletion, "should have inner agent completion")
	assert.True(t, foundOuterCompletion, "should have outer agent completion")
}

// consumeUntilInterrupt consumes events from the iterator until an interrupt is found or it's exhausted.
func consumeUntilInterrupt(iter *AsyncIterator[*AgentEvent]) (normalEvents []*AgentEvent, interruptEvent *AgentEvent) {
	for {
		event, ok := iter.Next()
		if !ok {
			break
		}
		if event.Action != nil && event.Action.Interrupted != nil {
			interruptEvent = event
			continue
		}
		normalEvents = append(normalEvents, event)
	}
	return
}

type returnDirectlyTool struct {
	name string
}

func (t *returnDirectlyTool) Info(_ context.Context) (*schema.ToolInfo, error) {
	return &schema.ToolInfo{
		Name: t.name,
		Desc: "A tool that returns directly",
	}, nil
}

func (t *returnDirectlyTool) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) {
	return "return directly result", nil
}

type interruptingTool struct {
	name string
}

func (i *interruptingTool) Info(_ context.Context) (*schema.ToolInfo, error) {
	return &schema.ToolInfo{
		Name: i.name,
		Desc: "A tool that interrupts",
	}, nil
}

func (i *interruptingTool) InvokableRun(ctx context.Context, _ string, _ ...tool.Option) (string, error) {
	if wasInterrupted, _, _ := compose.GetInterruptState[any](ctx); !wasInterrupted {
		return "", compose.Interrupt(ctx, "interrupt data")
	}

	if isResumeFlow, hasResumeData, data := compose.GetResumeContext[string](ctx); isResumeFlow && hasResumeData {
		return data, nil
	}

	return "resumed without data", nil
}

type twoToolCallModel struct {
	returnDirectlyToolName string
	interruptingToolName   string
	callCount              int
	receivedTools          []*schema.ToolInfo
	mu                     sync.Mutex
}

func (m *twoToolCallModel) Generate(_ context.Context, _ []*schema.Message, opts ...model.Option) (*schema.Message, error) {
	m.mu.Lock()
	m.callCount++
	callNum := m.callCount
	options := model.GetCommonOptions(&model.Options{}, opts...)
	if options.Tools != nil {
		m.receivedTools = options.Tools
	}
	m.mu.Unlock()

	if callNum == 1 {
		return &schema.Message{
			Role:    schema.Assistant,
			Content: "",
			ToolCalls: []schema.ToolCall{
				{
					ID:   "call_return_directly",
					Type: "function",
					Function: schema.FunctionCall{
						Name:      m.returnDirectlyToolName,
						Arguments: "{}",
					},
				},
				{
					ID:   "call_interrupting",
					Type: "function",
					Function: schema.FunctionCall{
						Name:      m.interruptingToolName,
						Arguments: "{}",
					},
				},
			},
		}, nil
	}
	return schema.AssistantMessage("final response", nil), nil
}

func (m *twoToolCallModel) Stream(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
	panic("not implemented")
}

func (m *twoToolCallModel) WithTools(_ []*schema.ToolInfo) (model.ToolCallingChatModel, error) {
	return m, nil
}

func (m *twoToolCallModel) GetReceivedTools() []*schema.ToolInfo {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.receivedTools
}

type dynamicTool struct {
	name string
}

func (t *dynamicTool) Info(_ context.Context) (*schema.ToolInfo, error) {
	return &schema.ToolInfo{
		Name: t.name,
		Desc: "A dynamically added tool",
	}, nil
}

func (t *dynamicTool) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) {
	return "dynamic tool result", nil
}

func TestReturnDirectlyEventSentAfterResume(t *testing.T) {
	ctx := context.Background()

	returnDirectlyToolName := "return_directly_tool"
	interruptingToolName := "interrupting_tool"
	dynamicToolName := "dynamic_tool"

	mdl := &twoToolCallModel{
		returnDirectlyToolName: returnDirectlyToolName,
		interruptingToolName:   interruptingToolName,
	}

	agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
		Name:        "TestAgent",
		Description: "Test agent for return directly + interrupt",
		Model:       mdl,
		ToolsConfig: ToolsConfig{
			ToolsNodeConfig: compose.ToolsNodeConfig{
				Tools: []tool.BaseTool{
					&returnDirectlyTool{name: returnDirectlyToolName},
					&interruptingTool{name: interruptingToolName},
				},
			},
			ReturnDirectly: map[string]bool{
				returnDirectlyToolName: true,
			},
		},
		Handlers: []ChatModelAgentMiddleware{
			&interruptTestToolsHandler{tools: []tool.BaseTool{&dynamicTool{name: dynamicToolName}}},
		},
	})
	assert.NoError(t, err)

	store := newMyStore()
	runner := NewRunner(ctx, RunnerConfig{
		Agent:           agent,
		EnableStreaming: false,
		CheckPointStore: store,
	})

	iter := runner.Query(ctx, "test input", WithCheckPointID("test_checkpoint"))

	var interruptEvent *AgentEvent
	for {
		event, ok := iter.Next()
		if !ok {
			break
		}
		if event.Action != nil && event.Action.Interrupted != nil {
			interruptEvent = event
		}
	}

	assert.NotNil(t, interruptEvent, "Should have an interrupt event")
	assert.NotEmpty(t, interruptEvent.Action.Interrupted.InterruptContexts)

	receivedToolsBeforeResume := mdl.GetReceivedTools()
	var hasDynamicToolBeforeResume bool
	for _, ti := range receivedToolsBeforeResume {
		if ti.Name == dynamicToolName {
			hasDynamicToolBeforeResume = true
		}
	}
	assert.True(t, hasDynamicToolBeforeResume, "Dynamic tool should be in tool list before interrupt")

	interruptID := interruptEvent.Action.Interrupted.InterruptContexts[0].ID
	resumeIter, err := runner.ResumeWithParams(ctx, "test_checkpoint", &ResumeParams{
		Targets: map[string]any{
			interruptID: "resume data",
		},
	})
	assert.NoError(t, err)

	var resumeEvents []*AgentEvent
	for {
		event, ok := resumeIter.Next()
		if !ok {
			break
		}
		resumeEvents = append(resumeEvents, event)
	}

	var hasReturnDirectlyEvent bool
	for _, e := range resumeEvents {
		if e.Output != nil && e.Output.MessageOutput != nil {
			if e.Output.MessageOutput.Role == schema.Tool && e.Output.MessageOutput.ToolName == returnDirectlyToolName {
				hasReturnDirectlyEvent = true
			}
		}
	}
	assert.True(t, hasReturnDirectlyEvent, "ReturnDirectlyEvent should be sent after resume")

	receivedToolsAfterResume := mdl.GetReceivedTools()
	var hasDynamicToolAfterResume bool
	for _, ti := range receivedToolsAfterResume {
		if ti.Name == dynamicToolName {
			hasDynamicToolAfterResume = true
		}
	}
	assert.True(t, hasDynamicToolAfterResume, "Dynamic tool should be in tool list after resume (bc.toolUpdated path)")
}

// streamErrorThenToolCallModel simulates a model that:
// - On the first Stream call: emits several good chunks then an error (triggering retry)
// - On the second Stream call (retry): returns a tool call message (success)
type streamErrorThenToolCallModel struct {
	callCount    int32
	toolCallName string
}

func (m *streamErrorThenToolCallModel) Generate(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
	return schema.AssistantMessage("final answer", nil), nil
}

func (m *streamErrorThenToolCallModel) Stream(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
	count := atomic.AddInt32(&m.callCount, 1)

	sr, sw := schema.Pipe[*schema.Message](10)
	go func() {
		defer sw.Close()
		if count == 1 {
			// First call: emit good chunks then error
			sw.Send(schema.AssistantMessage("chunk1", nil), nil)
			sw.Send(schema.AssistantMessage("chunk2", nil), nil)
			sw.Send(schema.AssistantMessage("chunk3", nil), nil)
			sw.Send(nil, errRetryAble)
			return
		}
		// Second call (retry): return tool call
		sw.Send(schema.AssistantMessage("", []schema.ToolCall{{
			ID:       "call-1",
			Function: schema.FunctionCall{Name: m.toolCallName, Arguments: "{}"},
		}}), nil)
	}()
	return sr, nil
}

func (m *streamErrorThenToolCallModel) WithTools(_ []*schema.ToolInfo) (model.ToolCallingChatModel, error) {
	return m, nil
}

// TestStreamRetryThenToolInterruptCheckpoint reproduces a bug where:
//  1. ChatModelAgent with ModelRetryConfig.MaxRetries = 2
//  2. First model Stream call emits good chunks then a retryable error
//  3. Retry succeeds, model returns a tool call
//  4. The tool triggers an interrupt, causing the Runner to save a checkpoint
//  5. The checkpoint save fails because the first (failed) model call's stream event
//     is in the session, and when MessageVariant.GobEncode consumes the stream,
//     it hits the error chunk and returns an encoding error.
func TestStreamRetryThenToolInterruptCheckpoint(t *testing.T) {
	ctx := context.Background()

	interruptToolName := "interrupt_tool"
	mdl := &streamErrorThenToolCallModel{toolCallName: interruptToolName}

	interruptTool := &interruptingTool{name: interruptToolName}

	agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
		Name:        "RetryInterruptAgent",
		Description: "Agent that retries model then tool interrupts",
		Instruction: "You are a test agent.",
		Model:       mdl,
		ModelRetryConfig: &ModelRetryConfig{
			MaxRetries:  2,
			IsRetryAble: func(ctx context.Context, err error) bool { return errors.Is(err, errRetryAble) },
		},
		ToolsConfig: ToolsConfig{
			ToolsNodeConfig: compose.ToolsNodeConfig{
				Tools: []tool.BaseTool{interruptTool},
			},
		},
	})
	assert.NoError(t, err)

	store := newMyStore()
	runner := NewRunner(ctx, RunnerConfig{
		Agent:           agent,
		EnableStreaming: true,
		CheckPointStore: store,
	})

	iter := runner.Run(ctx, []Message{schema.UserMessage("test query")}, WithCheckPointID("retry_interrupt_ckpt"))

	var events []*AgentEvent
	var checkpointErr error
	for {
		event, ok := iter.Next()
		if !ok {
			break
		}
		events = append(events, event)
		if event.Err != nil {
			checkpointErr = event.Err
			t.Logf("event error: %v", event.Err)
		}
	}

	// The bug: checkpoint save fails because the failed stream's error chunk
	// is encountered during gob encoding of the session events.
	// If the bug is fixed, checkpointErr should be nil and we should see an interrupt event.
	assert.NoError(t, checkpointErr, "checkpoint save should not fail due to failed stream's error in session")

	var hasInterrupt bool
	for _, event := range events {
		if event.Action != nil && event.Action.Interrupted != nil {
			hasInterrupt = true
		}
	}
	assert.True(t, hasInterrupt, "should receive an interrupt event from the tool")

	// Verify the model was called twice (first call errored, second succeeded)
	assert.Equal(t, int32(2), atomic.LoadInt32(&mdl.callCount), "model should be called exactly twice (1 failure + 1 retry)")
}


//go:build testing

package agent

import (
	"os"
	"strings"
	"testing"
	"time"

	"github.com/henrygd/beszel/internal/entities/system"
	"github.com/shirou/gopsutil/v4/disk"
	"github.com/stretchr/testify/assert"
)

func TestParseFilesystemEntry(t *testing.T) {
	tests := []struct {
		name         string
		input        string
		expectedFs   string
		expectedName string
	}{
		{
			name:         "simple device name",
			input:        "sda1",
			expectedFs:   "sda1",
			expectedName: "",
		},
		{
			name:         "device with custom name",
			input:        "sda1__my-storage",
			expectedFs:   "sda1",
			expectedName: "my-storage",
		},
		{
			name:         "full device path with custom name",
			input:        "/dev/sdb1__backup-drive",
			expectedFs:   "/dev/sdb1",
			expectedName: "backup-drive",
		},
		{
			name:         "NVMe device with custom name",
			input:        "nvme0n1p2__fast-ssd",
			expectedFs:   "nvme0n1p2",
			expectedName: "fast-ssd",
		},
		{
			name:         "whitespace trimmed",
			input:        "  sda2__trimmed-name  ",
			expectedFs:   "sda2",
			expectedName: "trimmed-name",
		},
		{
			name:         "empty custom name",
			input:        "sda3__",
			expectedFs:   "sda3",
			expectedName: "",
		},
		{
			name:         "empty device name",
			input:        "__just-custom",
			expectedFs:   "",
			expectedName: "just-custom",
		},
		{
			name:         "multiple underscores in custom name",
			input:        "sda1__my_custom_drive",
			expectedFs:   "sda1",
			expectedName: "my_custom_drive",
		},
		{
			name:         "custom name with spaces",
			input:        "sda1__My Storage Drive",
			expectedFs:   "sda1",
			expectedName: "My Storage Drive",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			fsEntry := strings.TrimSpace(tt.input)
			var fs, customName string
			if parts := strings.SplitN(fsEntry, "__", 2); len(parts) == 2 {
				fs = strings.TrimSpace(parts[0])
				customName = strings.TrimSpace(parts[1])
			} else {
				fs = fsEntry
			}

			assert.Equal(t, tt.expectedFs, fs)
			assert.Equal(t, tt.expectedName, customName)
		})
	}
}

func TestExtraFilesystemPartitionInfo(t *testing.T) {
	t.Run("uses partition device for label-only mountpoint", func(t *testing.T) {
		device, customName := extraFilesystemPartitionInfo(disk.PartitionStat{
			Device:     "/dev/sdc",
			Mountpoint: "/extra-filesystems/Share",
		})

		assert.Equal(t, "/dev/sdc", device)
		assert.Equal(t, "", customName)
	})

	t.Run("uses custom name from mountpoint suffix", func(t *testing.T) {
		device, customName := extraFilesystemPartitionInfo(disk.PartitionStat{
			Device:     "/dev/sdc",
			Mountpoint: "/extra-filesystems/sdc__Share",
		})

		assert.Equal(t, "/dev/sdc", device)
		assert.Equal(t, "Share", customName)
	})

	t.Run("falls back to folder device when partition device is unavailable", func(t *testing.T) {
		device, customName := extraFilesystemPartitionInfo(disk.PartitionStat{
			Mountpoint: "/extra-filesystems/sdc__Share",
		})

		assert.Equal(t, "sdc", device)
		assert.Equal(t, "Share", customName)
	})

	t.Run("supports custom name without folder device prefix", func(t *testing.T) {
		device, customName := extraFilesystemPartitionInfo(disk.PartitionStat{
			Device:     "/dev/sdc",
			Mountpoint: "/extra-filesystems/__Share",
		})

		assert.Equal(t, "/dev/sdc", device)
		assert.Equal(t, "Share", customName)
	})
}

func TestBuildFsStatRegistration(t *testing.T) {
	t.Run("uses basename for non-windows exact io match", func(t *testing.T) {
		key, stats, ok := registerFilesystemStats(
			map[string]*system.FsStats{},
			"/dev/sda1",
			"/mnt/data",
			false,
			"archive",
			fsRegistrationContext{
				isWindows: false,
				diskIoCounters: map[string]disk.IOCountersStat{
					"sda1": {Name: "sda1"},
				},
			},
		)

		assert.True(t, ok)
		assert.Equal(t, "sda1", key)
		assert.Equal(t, "/mnt/data", stats.Mountpoint)
		assert.Equal(t, "archive", stats.Name)
		assert.False(t, stats.Root)
	})

	t.Run("maps root partition to io device by prefix", func(t *testing.T) {
		key, stats, ok := registerFilesystemStats(
			map[string]*system.FsStats{},
			"/dev/ada0p2",
			"/",
			true,
			"",
			fsRegistrationContext{
				isWindows: false,
				diskIoCounters: map[string]disk.IOCountersStat{
					"ada0": {Name: "ada0", ReadBytes: 1000, WriteBytes: 1000},
				},
			},
		)

		assert.True(t, ok)
		assert.Equal(t, "ada0", key)
		assert.True(t, stats.Root)
		assert.Equal(t, "/", stats.Mountpoint)
	})

	t.Run("uses filesystem setting as root fallback", func(t *testing.T) {
		key, _, ok := registerFilesystemStats(
			map[string]*system.FsStats{},
			"overlay",
			"/",
			true,
			"",
			fsRegistrationContext{
				filesystem: "nvme0n1p2",
				isWindows:  false,
				diskIoCounters: map[string]disk.IOCountersStat{
					"nvme0n1": {Name: "nvme0n1", ReadBytes: 1000, WriteBytes: 1000},
				},
			},
		)

		assert.True(t, ok)
		assert.Equal(t, "nvme0n1", key)
	})

	t.Run("prefers parsed extra-filesystems device over mapper device", func(t *testing.T) {
		key, stats, ok := registerFilesystemStats(
			map[string]*system.FsStats{},
			"/dev/mapper/luks-2bcb02be-999d-4417-8d18-5c61e660fb6e",
			"/extra-filesystems/nvme0n1p2__Archive",
			false,
			"Archive",
			fsRegistrationContext{
				isWindows: false,
				diskIoCounters: map[string]disk.IOCountersStat{
					"dm-1":      {Name: "dm-1", Label: "luks-2bcb02be-999d-4417-8d18-5c61e660fb6e"},
					"nvme0n1p2": {Name: "nvme0n1p2"},
				},
			},
		)

		assert.True(t, ok)
		assert.Equal(t, "nvme0n1p2", key)
		assert.Equal(t, "Archive", stats.Name)
	})

	t.Run("falls back to mapper io device when folder device cannot be resolved", func(t *testing.T) {
		key, stats, ok := registerFilesystemStats(
			map[string]*system.FsStats{},
			"/dev/mapper/luks-2bcb02be-999d-4417-8d18-5c61e660fb6e",
			"/extra-filesystems/Archive",
			false,
			"Archive",
			fsRegistrationContext{
				isWindows: false,
				diskIoCounters: map[string]disk.IOCountersStat{
					"dm-1": {Name: "dm-1", Label: "luks-2bcb02be-999d-4417-8d18-5c61e660fb6e"},
				},
			},
		)

		assert.True(t, ok)
		assert.Equal(t, "dm-1", key)
		assert.Equal(t, "Archive", stats.Name)
	})

	t.Run("uses full device name on windows", func(t *testing.T) {
		key, _, ok := registerFilesystemStats(
			map[string]*system.FsStats{},
			`C:`,
			`C:\\`,
			false,
			"",
			fsRegistrationContext{
				isWindows: true,
				diskIoCounters: map[string]disk.IOCountersStat{
					`C:`: {Name: `C:`},
				},
			},
		)

		assert.True(t, ok)
		assert.Equal(t, `C:`, key)
	})

	t.Run("skips existing key", func(t *testing.T) {
		key, stats, ok := registerFilesystemStats(
			map[string]*system.FsStats{"sda1": {Mountpoint: "/existing"}},
			"/dev/sda1",
			"/mnt/data",
			false,
			"",
			fsRegistrationContext{
				isWindows: false,
				diskIoCounters: map[string]disk.IOCountersStat{
					"sda1": {Name: "sda1"},
				},
			},
		)

		assert.False(t, ok)
		assert.Empty(t, key)
		assert.Nil(t, stats)
	})
}

func TestAddConfiguredRootFs(t *testing.T) {
	t.Run("adds root from matching partition", func(t *testing.T) {
		agent := &Agent{fsStats: make(map[string]*system.FsStats)}
		discovery := diskDiscovery{
			agent:          agent,
			rootMountPoint: "/",
			partitions:     []disk.PartitionStat{{Device: "/dev/ada0p2", Mountpoint: "/"}},
			ctx: fsRegistrationContext{
				filesystem: "/dev/ada0p2",
				isWindows:  false,
				diskIoCounters: map[string]disk.IOCountersStat{
					"ada0": {Name: "ada0", ReadBytes: 1000, WriteBytes: 1000},
				},
			},
		}

		ok := discovery.addConfiguredRootFs()

		assert.True(t, ok)
		stats, exists := agent.fsStats["ada0"]
		assert.True(t, exists)
		assert.True(t, stats.Root)
		assert.Equal(t, "/", stats.Mountpoint)
	})

	t.Run("adds root from io device when partition is missing", func(t *testing.T) {
		agent := &Agent{fsStats: make(map[string]*system.FsStats)}
		discovery := diskDiscovery{
			agent:          agent,
			rootMountPoint: "/sysroot",
			ctx: fsRegistrationContext{
				filesystem: "zroot",
				isWindows:  false,
				diskIoCounters: map[string]disk.IOCountersStat{
					"nda0": {Name: "nda0", Label: "zroot", ReadBytes: 1000, WriteBytes: 1000},
				},
			},
		}

		ok := discovery.addConfiguredRootFs()

		assert.True(t, ok)
		stats, exists := agent.fsStats["nda0"]
		assert.True(t, exists)
		assert.True(t, stats.Root)
		assert.Equal(t, "/sysroot", stats.Mountpoint)
	})

	t.Run("returns false when filesystem cannot be resolved", func(t *testing.T) {
		agent := &Agent{fsStats: make(map[string]*system.FsStats)}
		discovery := diskDiscovery{
			agent:          agent,
			rootMountPoint: "/",
			ctx: fsRegistrationContext{
				filesystem:     "missing-disk",
				isWindows:      false,
				diskIoCounters: map[string]disk.IOCountersStat{},
			},
		}

		ok := discovery.addConfiguredRootFs()

		assert.False(t, ok)
		assert.Empty(t, agent.fsStats)
	})
}

func TestAddPartitionRootFs(t *testing.T) {
	t.Run("adds root from fallback partition candidate", func(t *testing.T) {
		agent := &Agent{fsStats: make(map[string]*system.FsStats)}
		discovery := diskDiscovery{
			agent: agent,
			ctx: fsRegistrationContext{
				isWindows: false,
				diskIoCounters: map[string]disk.IOCountersStat{
					"nvme0n1": {Name: "nvme0n1", ReadBytes: 1000, WriteBytes: 1000},
				},
			},
		}

		ok := discovery.addPartitionRootFs("/dev/nvme0n1p2", "/")

		assert.True(t, ok)
		stats, exists := agent.fsStats["nvme0n1"]
		assert.True(t, exists)
		assert.True(t, stats.Root)
		assert.Equal(t, "/", stats.Mountpoint)
	})

	t.Run("returns false when no io device matches", func(t *testing.T) {
		agent := &Agent{fsStats: make(map[string]*system.FsStats)}
		discovery := diskDiscovery{agent: agent, ctx: fsRegistrationContext{diskIoCounters: map[string]disk.IOCountersStat{}}}

		ok := discovery.addPartitionRootFs("/dev/mapper/root", "/")

		assert.False(t, ok)
		assert.Empty(t, agent.fsStats)
	})
}

func TestAddLastResortRootFs(t *testing.T) {
	t.Run("uses most active io device when available", func(t *testing.T) {
		agent := &Agent{fsStats: make(map[string]*system.FsStats)}
		discovery := diskDiscovery{agent: agent, rootMountPoint: "/", ctx: fsRegistrationContext{diskIoCounters: map[string]disk.IOCountersStat{
			"sda": {Name: "sda", ReadBytes: 5000, WriteBytes: 5000},
			"sdb": {Name: "sdb", ReadBytes: 1000, WriteBytes: 1000},
		}}}

		discovery.addLastResortRootFs()

		stats, exists := agent.fsStats["sda"]
		assert.True(t, exists)
		assert.True(t, stats.Root)
	})

	t.Run("falls back to root key when mountpoint basename collides", func(t *testing.T) {
		agent := &Agent{fsStats: map[string]*system.FsStats{
			"sysroot": {Mountpoint: "/extra-filesystems/sysroot"},
		}}
		discovery := diskDiscovery{agent: agent, rootMountPoint: "/sysroot", ctx: fsRegistrationContext{diskIoCounters: map[string]disk.IOCountersStat{}}}

		discovery.addLastResortRootFs()

		stats, exists := agent.fsStats["root"]
		assert.True(t, exists)
		assert.True(t, stats.Root)
		assert.Equal(t, "/sysroot", stats.Mountpoint)
	})
}

func TestAddConfiguredExtraFsEntry(t *testing.T) {
	t.Run("uses matching partition when present", func(t *testing.T) {
		agent := &Agent{fsStats: make(map[string]*system.FsStats)}
		discovery := diskDiscovery{
			agent:      agent,
			partitions: []disk.PartitionStat{{Device: "/dev/sdb1", Mountpoint: "/mnt/backup"}},
			usageFn: func(string) (*disk.UsageStat, error) {
				t.Fatal("usage fallback should not be called when partition matches")
				return nil, nil
			},
			ctx: fsRegistrationContext{
				isWindows: false,
				diskIoCounters: map[string]disk.IOCountersStat{
					"sdb1": {Name: "sdb1"},
				},
			},
		}

		discovery.addConfiguredExtraFsEntry("sdb1", "backup")

		stats, exists := agent.fsStats["sdb1"]
		assert.True(t, exists)
		assert.Equal(t, "/mnt/backup", stats.Mountpoint)
		assert.Equal(t, "backup", stats.Name)
	})

	t.Run("falls back to usage-validated path", func(t *testing.T) {
		agent := &Agent{fsStats: make(map[string]*system.FsStats)}
		discovery := diskDiscovery{
			agent: agent,
			usageFn: func(path string) (*disk.UsageStat, error) {
				assert.Equal(t, "/srv/archive", path)
				return &disk.UsageStat{}, nil
			},
			ctx: fsRegistrationContext{
				isWindows: false,
				diskIoCounters: map[string]disk.IOCountersStat{
					"archive": {Name: "archive"},
				},
			},
		}

		discovery.addConfiguredExtraFsEntry("/srv/archive", "archive")

		stats, exists := agent.fsStats["archive"]
		assert.True(t, exists)
		assert.Equal(t, "/srv/archive", stats.Mountpoint)
		assert.Equal(t, "archive", stats.Name)
	})

	t.Run("ignores invalid filesystem entry", func(t *testing.T) {
		agent := &Agent{fsStats: make(map[string]*system.FsStats)}
		discovery := diskDiscovery{
			agent: agent,
			usageFn: func(string) (*disk.UsageStat, error) {
				return nil, os.ErrNotExist
			},
		}

		discovery.addConfiguredExtraFsEntry("/missing/archive", "")

		assert.Empty(t, agent.fsStats)
	})
}

func TestAddConfiguredExtraFilesystems(t *testing.T) {
	t.Run("parses and registers multiple configured filesystems", func(t *testing.T) {
		agent := &Agent{fsStats: make(map[string]*system.FsStats)}
		discovery := diskDiscovery{
			agent:      agent,
			partitions: []disk.PartitionStat{{Device: "/dev/sda1", Mountpoint: "/mnt/fast"}},
			usageFn: func(path string) (*disk.UsageStat, error) {
				if path == "/srv/archive" {
					return &disk.UsageStat{}, nil
				}
				return nil, os.ErrNotExist
			},
			ctx: fsRegistrationContext{
				isWindows: false,
				diskIoCounters: map[string]disk.IOCountersStat{
					"sda1":    {Name: "sda1"},
					"archive": {Name: "archive"},
				},
			},
		}

		discovery.addConfiguredExtraFilesystems("sda1__fast,/srv/archive__cold")

		assert.Contains(t, agent.fsStats, "sda1")
		assert.Equal(t, "fast", agent.fsStats["sda1"].Name)
		assert.Contains(t, agent.fsStats, "archive")
		assert.Equal(t, "cold", agent.fsStats["archive"].Name)
	})
}

func TestAddExtraFilesystemFolders(t *testing.T) {
	t.Run("adds missing folders and skips existing mountpoints", func(t *testing.T) {
		agent := &Agent{fsStats: map[string]*system.FsStats{
			"existing": {Mountpoint: "/extra-filesystems/existing"},
		}}
		discovery := diskDiscovery{
			agent: agent,
			ctx: fsRegistrationContext{
				isWindows: false,
				efPath:    "/extra-filesystems",
				diskIoCounters: map[string]disk.IOCountersStat{
					"newdisk": {Name: "newdisk"},
				},
			},
		}

		discovery.addExtraFilesystemFolders([]string{"existing", "newdisk__Archive"})

		assert.Len(t, agent.fsStats, 2)
		stats, exists := agent.fsStats["newdisk"]
		assert.True(t, exists)
		assert.Equal(t, "/extra-filesystems/newdisk__Archive", stats.Mountpoint)
		assert.Equal(t, "Archive", stats.Name)
	})
}

func TestAddPartitionExtraFs(t *testing.T) {
	makeDiscovery := func(agent *Agent) diskDiscovery {
		return diskDiscovery{
			agent: agent,
			ctx: fsRegistrationContext{
				isWindows: false,
				efPath:    "/extra-filesystems",
				diskIoCounters: map[string]disk.IOCountersStat{
					"nvme0n1p1": {Name: "nvme0n1p1"},
					"nvme1n1":   {Name: "nvme1n1"},
				},
			},
		}
	}

	t.Run("registers direct child of extra-filesystems", func(t *testing.T) {
		agent := &Agent{fsStats: make(map[string]*system.FsStats)}
		d := makeDiscovery(agent)

		d.addPartitionExtraFs(disk.PartitionStat{
			Device:     "/dev/nvme0n1p1",
			Mountpoint: "/extra-filesystems/nvme0n1p1__caddy1-root",
		})

		stats, exists := agent.fsStats["nvme0n1p1"]
		assert.True(t, exists)
		assert.Equal(t, "/extra-filesystems/nvme0n1p1__caddy1-root", stats.Mountpoint)
		assert.Equal(t, "caddy1-root", stats.Name)
	})

	t.Run("skips nested mount under extra-filesystem bind mount", func(t *testing.T) {
		agent := &Agent{fsStats: make(map[string]*system.FsStats)}
		d := makeDiscovery(agent)

		// These simulate the virtual mounts that appear when host / is bind-mounted
		// with disk.Partitions(all=true) — e.g. /proc, /sys, /dev visible under the mount.
		for _, nested := range []string{
			"/extra-filesystems/nvme0n1p1__caddy1-root/proc",
			"/extra-filesystems/nvme0n1p1__caddy1-root/sys",
			"/extra-filesystems/nvme0n1p1__caddy1-root/dev",
			"/extra-filesystems/nvme0n1p1__caddy1-root/run",
		} {
			d.addPartitionExtraFs(disk.PartitionStat{Device: "tmpfs", Mountpoint: nested})
		}

		assert.Empty(t, agent.fsStats)
	})

	t.Run("registers both direct children, skips their nested mounts", func(t *testing.T) {
		agent := &Agent{fsStats: make(map[string]*system.FsStats)}
		d := makeDiscovery(agent)

		partitions := []disk.PartitionStat{
			{Device: "/dev/nvme0n1p1", Mountpoint: "/extra-filesystems/nvme0n1p1__caddy1-root"},
			{Device: "/dev/nvme1n1", Mountpoint: "/extra-filesystems/nvme1n1__caddy1-docker"},
			{Device: "proc", Mountpoint: "/extra-filesystems/nvme0n1p1__caddy1-root/proc"},
			{Device: "sysfs", Mountpoint: "/extra-filesystems/nvme0n1p1__caddy1-root/sys"},
			{Device: "overlay", Mountpoint: "/extra-filesystems/nvme0n1p1__caddy1-root/var/lib/docker"},
		}
		for _, p := range partitions {
			d.addPartitionExtraFs(p)
		}

		assert.Len(t, agent.fsStats, 2)
		assert.Equal(t, "caddy1-root", agent.fsStats["nvme0n1p1"].Name)
		assert.Equal(t, "caddy1-docker", agent.fsStats["nvme1n1"].Name)
	})

	t.Run("skips partition not under extra-filesystems", func(t *testing.T) {
		agent := &Agent{fsStats: make(map[string]*system.FsStats)}
		d := makeDiscovery(agent)

		d.addPartitionExtraFs(disk.PartitionStat{
			Device:     "/dev/nvme0n1p1",
			Mountpoint: "/",
		})

		assert.Empty(t, agent.fsStats)
	})
}

func TestFindIoDevice(t *testing.T) {
	t.Run("matches by device name", func(t *testing.T) {
		ioCounters := map[string]disk.IOCountersStat{
			"sda": {Name: "sda"},
			"sdb": {Name: "sdb"},
		}

		device, ok := findIoDevice("sdb", ioCounters)
		assert.True(t, ok)
		assert.Equal(t, "sdb", device)
	})

	t.Run("matches by device label", func(t *testing.T) {
		ioCounters := map[string]disk.IOCountersStat{
			"sda": {Name: "sda", Label: "rootfs"},
			"sdb": {Name: "sdb"},
		}

		device, ok := findIoDevice("rootfs", ioCounters)
		assert.True(t, ok)
		assert.Equal(t, "sda", device)
	})

	t.Run("returns no match when not found", func(t *testing.T) {
		ioCounters := map[string]disk.IOCountersStat{
			"sda": {Name: "sda"},
			"sdb": {Name: "sdb"},
		}

		device, ok := findIoDevice("nvme0n1p1", ioCounters)
		assert.False(t, ok)
		assert.Equal(t, "", device)
	})

	t.Run("uses uncertain unique prefix fallback", func(t *testing.T) {
		ioCounters := map[string]disk.IOCountersStat{
			"nvme0n1": {Name: "nvme0n1"},
			"sda":     {Name: "sda"},
		}

		device, ok := findIoDevice("nvme0n1p2", ioCounters)
		assert.True(t, ok)
		assert.Equal(t, "nvme0n1", device)
	})

	t.Run("uses dominant activity when prefix matches are ambiguous", func(t *testing.T) {
		ioCounters := map[string]disk.IOCountersStat{
			"sda": {Name: "sda", ReadBytes: 5000, WriteBytes: 5000, ReadCount: 100, WriteCount: 100},
			"sdb": {Name: "sdb", ReadBytes: 1000, WriteBytes: 1000, ReadCount: 50, WriteCount: 50},
		}

		device, ok := findIoDevice("sd", ioCounters)
		assert.True(t, ok)
		assert.Equal(t, "sda", device)
	})

	t.Run("uses highest activity when ambiguous without dominance", func(t *testing.T) {
		ioCounters := map[string]disk.IOCountersStat{
			"sda": {Name: "sda", ReadBytes: 3000, WriteBytes: 3000, ReadCount: 50, WriteCount: 50},
			"sdb": {Name: "sdb", ReadBytes: 2500, WriteBytes: 2500, ReadCount: 40, WriteCount: 40},
		}

		device, ok := findIoDevice("sd", ioCounters)
		assert.True(t, ok)
		assert.Equal(t, "sda", device)
	})

	t.Run("matches /dev/-prefixed partition to parent disk", func(t *testing.T) {
		ioCounters := map[string]disk.IOCountersStat{
			"nda0": {Name: "nda0", ReadBytes: 1000, WriteBytes: 1000},
		}

		device, ok := findIoDevice("/dev/nda0p2", ioCounters)
		assert.True(t, ok)
		assert.Equal(t, "nda0", device)
	})

	t.Run("uses deterministic name tie-breaker", func(t *testing.T) {
		ioCounters := map[string]disk.IOCountersStat{
			"sdb": {Name: "sdb", ReadBytes: 2000, WriteBytes: 2000, ReadCount: 10, WriteCount: 10},
			"sda": {Name: "sda", ReadBytes: 2000, WriteBytes: 2000, ReadCount: 10, WriteCount: 10},
		}

		device, ok := findIoDevice("sd", ioCounters)
		assert.True(t, ok)
		assert.Equal(t, "sda", device)
	})
}

func TestFilesystemMatchesPartitionSetting(t *testing.T) {
	p := disk.PartitionStat{Device: "/dev/ada0p2", Mountpoint: "/"}

	t.Run("matches mountpoint setting", func(t *testing.T) {
		assert.True(t, filesystemMatchesPartitionSetting("/", p))
	})

	t.Run("matches exact partition setting", func(t *testing.T) {
		assert.True(t, filesystemMatchesPartitionSetting("ada0p2", p))
		assert.True(t, filesystemMatchesPartitionSetting("/dev/ada0p2", p))
	})

	t.Run("matches prefix-style parent setting", func(t *testing.T) {
		assert.True(t, filesystemMatchesPartitionSetting("ada0", p))
		assert.True(t, filesystemMatchesPartitionSetting("/dev/ada0", p))
	})

	t.Run("does not match unrelated device", func(t *testing.T) {
		assert.False(t, filesystemMatchesPartitionSetting("sda", p))
		assert.False(t, filesystemMatchesPartitionSetting("nvme0n1", p))
		assert.False(t, filesystemMatchesPartitionSetting("", p))
	})
}

func TestMostActiveIoDevice(t *testing.T) {
	t.Run("returns most active device", func(t *testing.T) {
		ioCounters := map[string]disk.IOCountersStat{
			"nda0": {Name: "nda0", ReadBytes: 5000, WriteBytes: 5000, ReadCount: 100, WriteCount: 100},
			"nda1": {Name: "nda1", ReadBytes: 1000, WriteBytes: 1000, ReadCount: 50, WriteCount: 50},
		}
		assert.Equal(t, "nda0", mostActiveIoDevice(ioCounters))
	})

	t.Run("uses deterministic tie-breaker", func(t *testing.T) {
		ioCounters := map[string]disk.IOCountersStat{
			"sdb": {Name: "sdb", ReadBytes: 1000, WriteBytes: 1000, ReadCount: 10, WriteCount: 10},
			"sda": {Name: "sda", ReadBytes: 1000, WriteBytes: 1000, ReadCount: 10, WriteCount: 10},
		}
		assert.Equal(t, "sda", mostActiveIoDevice(ioCounters))
	})

	t.Run("returns empty for empty map", func(t *testing.T) {
		assert.Equal(t, "", mostActiveIoDevice(map[string]disk.IOCountersStat{}))
	})
}

func TestIsDockerSpecialMountpoint(t *testing.T) {
	testCases := []struct {
		name       string
		mountpoint string
		expected   bool
	}{
		{name: "hosts", mountpoint: "/etc/hosts", expected: true},
		{name: "resolv", mountpoint: "/etc/resolv.conf", expected: true},
		{name: "hostname", mountpoint: "/etc/hostname", expected: true},
		{name: "root", mountpoint: "/", expected: false},
		{name: "passwd", mountpoint: "/etc/passwd", expected: false},
		{name: "extra-filesystem", mountpoint: "/extra-filesystems/sda1", expected: false},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			assert.Equal(t, tc.expected, isDockerSpecialMountpoint(tc.mountpoint))
		})
	}
}

func TestInitializeDiskInfoWithCustomNames(t *testing.T) {
	// Test with custom names
	t.Setenv("EXTRA_FILESYSTEMS", "sda1__my-storage,/dev/sdb1__backup-drive,nvme0n1p2")

	// Mock disk partitions (we'll just test the parsing logic)
	// Since the actual disk operations are system-dependent, we'll focus on the parsing
	testCases := []struct {
		envValue      string
		expectedFs    []string
		expectedNames map[string]string
	}{
		{
			envValue:   "sda1__my-storage,sdb1__backup-drive",
			expectedFs: []string{"sda1", "sdb1"},
			expectedNames: map[string]string{
				"sda1": "my-storage",
				"sdb1": "backup-drive",
			},
		},
		{
			envValue:   "sda1,nvme0n1p2__fast-ssd",
			expectedFs: []string{"sda1", "nvme0n1p2"},
			expectedNames: map[string]string{
				"nvme0n1p2": "fast-ssd",
			},
		},
	}

	for _, tc := range testCases {
		t.Run("env_"+tc.envValue, func(t *testing.T) {
			t.Setenv("EXTRA_FILESYSTEMS", tc.envValue)

			// Create mock partitions that would match our test cases
			partitions := []disk.PartitionStat{}
			for _, fs := range tc.expectedFs {
				if strings.HasPrefix(fs, "/dev/") {
					partitions = append(partitions, disk.PartitionStat{
						Device:     fs,
						Mountpoint: fs,
					})
				} else {
					partitions = append(partitions, disk.PartitionStat{
						Device:     "/dev/" + fs,
						Mountpoint: "/" + fs,
					})
				}
			}

			// Test the parsing logic by calling the relevant part
			// We'll create a simplified version to test just the parsing
			extraFilesystems := tc.envValue
			for fsEntry := range strings.SplitSeq(extraFilesystems, ",") {
				// Parse the entry
				fsEntry = strings.TrimSpace(fsEntry)
				var fs, customName string
				if parts := strings.SplitN(fsEntry, "__", 2); len(parts) == 2 {
					fs = strings.TrimSpace(parts[0])
					customName = strings.TrimSpace(parts[1])
				} else {
					fs = fsEntry
				}

				// Verify the device is in our expected list
				assert.Contains(t, tc.expectedFs, fs, "parsed device should be in expected list")

				// Check if custom name should exist
				if expectedName, exists := tc.expectedNames[fs]; exists {
					assert.Equal(t, expectedName, customName, "custom name should match expected")
				} else {
					assert.Empty(t, customName, "custom name should be empty when not expected")
				}
			}
		})
	}
}

func TestFsStatsWithCustomNames(t *testing.T) {
	// Test that FsStats properly stores custom names
	fsStats := &system.FsStats{
		Mountpoint: "/mnt/storage",
		Name:       "my-custom-storage",
		DiskTotal:  100.0,
		DiskUsed:   50.0,
	}

	assert.Equal(t, "my-custom-storage", fsStats.Name)
	assert.Equal(t, "/mnt/storage", fsStats.Mountpoint)
	assert.Equal(t, 100.0, fsStats.DiskTotal)
	assert.Equal(t, 50.0, fsStats.DiskUsed)
}

func TestExtraFsKeyGeneration(t *testing.T) {
	// Test the logic for generating ExtraFs keys with custom names
	testCases := []struct {
		name        string
		deviceName  string
		customName  string
		expectedKey string
	}{
		{
			name:        "with custom name",
			deviceName:  "sda1",
			customName:  "my-storage",
			expectedKey: "my-storage",
		},
		{
			name:        "without custom name",
			deviceName:  "sda1",
			customName:  "",
			expectedKey: "sda1",
		},
		{
			name:        "empty custom name falls back to device",
			deviceName:  "nvme0n1p2",
			customName:  "",
			expectedKey: "nvme0n1p2",
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			// Simulate the key generation logic from agent.go
			key := tc.deviceName
			if tc.customName != "" {
				key = tc.customName
			}
			assert.Equal(t, tc.expectedKey, key)
		})
	}
}

func TestDiskUsageCaching(t *testing.T) {
	t.Run("caching disabled updates all filesystems", func(t *testing.T) {
		agent := &Agent{
			fsStats: map[string]*system.FsStats{
				"sda": {Root: true, Mountpoint: "/"},
				"sdb": {Root: false, Mountpoint: "/mnt/storage"},
			},
			diskUsageCacheDuration: 0, // caching disabled
		}

		var stats system.Stats
		agent.updateDiskUsage(&stats)

		// Both should be updated (non-zero values from disk.Usage)
		// Root stats should be populated in systemStats
		assert.True(t, agent.lastDiskUsageUpdate.IsZero() || !agent.lastDiskUsageUpdate.IsZero(),
			"lastDiskUsageUpdate should be set when caching is disabled")
	})

	t.Run("caching enabled always updates root filesystem", func(t *testing.T) {
		agent := &Agent{
			fsStats: map[string]*system.FsStats{
				"sda": {Root: true, Mountpoint: "/", DiskTotal: 100, DiskUsed: 50},
				"sdb": {Root: false, Mountpoint: "/mnt/storage", DiskTotal: 200, DiskUsed: 100},
			},
			diskUsageCacheDuration: 1 * time.Hour,
			lastDiskUsageUpdate:    time.Now(), // cache is fresh
		}

		// Store original extra fs values
		originalExtraTotal := agent.fsStats["sdb"].DiskTotal
		originalExtraUsed := agent.fsStats["sdb"].DiskUsed

		var stats system.Stats
		agent.updateDiskUsage(&stats)

		// Root should be updated (systemStats populated from disk.Usage call)
		// We can't easily check if disk.Usage was called, but we verify the flow works

		// Extra filesystem should retain cached values (not reset)
		assert.Equal(t, originalExtraTotal, agent.fsStats["sdb"].DiskTotal,
			"extra filesystem DiskTotal should be unchanged when cached")
		assert.Equal(t, originalExtraUsed, agent.fsStats["sdb"].DiskUsed,
			"extra filesystem DiskUsed should be unchanged when cached")
	})

	t.Run("first call always updates all filesystems", func(t *testing.T) {
		agent := &Agent{
			fsStats: map[string]*system.FsStats{
				"sda": {Root: true, Mountpoint: "/"},
				"sdb": {Root: false, Mountpoint: "/mnt/storage"},
			},
			diskUsageCacheDuration: 1 * time.Hour,
			// lastDiskUsageUpdate is zero (first call)
		}

		var stats system.Stats
		agent.updateDiskUsage(&stats)

		// After first call, lastDiskUsageUpdate should be set
		assert.False(t, agent.lastDiskUsageUpdate.IsZero(),
			"lastDiskUsageUpdate should be set after first call")
	})

	t.Run("expired cache updates extra filesystems", func(t *testing.T) {
		agent := &Agent{
			fsStats: map[string]*system.FsStats{
				"sda": {Root: true, Mountpoint: "/"},
				"sdb": {Root: false, Mountpoint: "/mnt/storage"},
			},
			diskUsageCacheDuration: 1 * time.Millisecond,
			lastDiskUsageUpdate:    time.Now().Add(-1 * time.Second), // cache expired
		}

		var stats system.Stats
		agent.updateDiskUsage(&stats)

		// lastDiskUsageUpdate should be refreshed since cache expired
		assert.True(t, time.Since(agent.lastDiskUsageUpdate) < time.Second,
			"lastDiskUsageUpdate should be refreshed when cache expires")
	})
}

func TestHasSameDiskUsage(t *testing.T) {
	const toleranceBytes uint64 = 16 * 1024 * 1024

	t.Run("returns true when totals and usage are equal", func(t *testing.T) {
		a := &disk.UsageStat{Total: 100 * 1024 * 1024 * 1024, Used: 42 * 1024 * 1024 * 1024}
		b := &disk.UsageStat{Total: 100 * 1024 * 1024 * 1024, Used: 42 * 1024 * 1024 * 1024}
		assert.True(t, hasSameDiskUsage(a, b))
	})

	t.Run("returns true within tolerance", func(t *testing.T) {
		a := &disk.UsageStat{Total: 100 * 1024 * 1024 * 1024, Used: 42 * 1024 * 1024 * 1024}
		b := &disk.UsageStat{
			Total: a.Total + toleranceBytes - 1,
			Used:  a.Used - toleranceBytes + 1,
		}
		assert.True(t, hasSameDiskUsage(a, b))
	})

	t.Run("returns false when total exceeds tolerance", func(t *testing.T) {
		a := &disk.UsageStat{Total: 100 * 1024 * 1024 * 1024, Used: 42 * 1024 * 1024 * 1024}
		b := &disk.UsageStat{
			Total: a.Total + toleranceBytes + 1,
			Used:  a.Used,
		}
		assert.False(t, hasSameDiskUsage(a, b))
	})

	t.Run("returns false for nil or zero total", func(t *testing.T) {
		assert.False(t, hasSameDiskUsage(nil, &disk.UsageStat{Total: 1, Used: 1}))
		assert.False(t, hasSameDiskUsage(&disk.UsageStat{Total: 1, Used: 1}, nil))
		assert.False(t, hasSameDiskUsage(&disk.UsageStat{Total: 0, Used: 0}, &disk.UsageStat{Total: 1, Used: 1}))
	})
}

func TestInitializeDiskIoStatsResetsTrackedDevices(t *testing.T) {
	agent := &Agent{
		fsStats: map[string]*system.FsStats{
			"sda": {},
			"sdb": {},
		},
		fsNames: []string{"stale", "sda"},
	}

	agent.initializeDiskIoStats(map[string]disk.IOCountersStat{
		"sda": {Name: "sda", ReadBytes: 10, WriteBytes: 20},
		"sdb": {Name: "sdb", ReadBytes: 30, WriteBytes: 40},
	})

	assert.ElementsMatch(t, []string{"sda", "sdb"}, agent.fsNames)
	assert.Len(t, agent.fsNames, 2)
	assert.Equal(t, uint64(10), agent.fsStats["sda"].TotalRead)
	assert.Equal(t, uint64(20), agent.fsStats["sda"].TotalWrite)
	assert.False(t, agent.fsStats["sda"].Time.IsZero())
	assert.False(t, agent.fsStats["sdb"].Time.IsZero())

	agent.initializeDiskIoStats(map[string]disk.IOCountersStat{
		"sdb": {Name: "sdb", ReadBytes: 50, WriteBytes: 60},
	})

	assert.Equal(t, []string{"sdb"}, agent.fsNames)
	assert.Equal(t, uint64(50), agent.fsStats["sdb"].TotalRead)
	assert.Equal(t, uint64(60), agent.fsStats["sdb"].TotalWrite)
}


"""
Custom HTTP Headers via a custom Watchdog.

Creates a custom watchdog that listens to TabCreatedEvent and injects
custom HTTP headers into every new tab using Network.setExtraHTTPHeaders.

Note: The CDP EventRegistry only supports one handler per event method,
so registering directly on Target.attachedToTarget would replace the
internal SessionManager handler.  Using the browser-use event system
(TabCreatedEvent) avoids this and fires after the target is fully set up.

Note: Network.setExtraHTTPHeaders is a full replacement (not additive).

Verified by navigating to https://httpbin.org/headers in a new tab.
"""

import asyncio
import os
import sys
from typing import ClassVar

sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

from bubus import BaseEvent
from dotenv import load_dotenv

load_dotenv()

from browser_use import Agent, Browser, ChatBrowserUse
from browser_use.browser.events import AgentFocusChangedEvent, TabCreatedEvent
from browser_use.browser.watchdog_base import BaseWatchdog

CUSTOM_HEADERS = {
	'X-Custom-Auth': 'Bearer my-secret-token',
	'X-Request-Source': 'browser-use-agent',
	'X-Trace-Id': 'example-trace-12345',
}

class CustomHeadersWatchdog(BaseWatchdog):
	"""Injects custom HTTP headers on every new tab and focus change.

	Listens to both TabCreatedEvent (new tabs) and AgentFocusChangedEvent
	(tab switches) because headers are bound to a CDP session, and sessions
	can be recreated on cross-origin navigations or tab switches.
	"""

	LISTENS_TO: ClassVar[list[type[BaseEvent]]] = [TabCreatedEvent, AgentFocusChangedEvent]
	EMITS: ClassVar[list[type[BaseEvent]]] = []

	async def on_TabCreatedEvent(self, event: TabCreatedEvent) -> None:
		"""Set extra headers when a new tab is created."""
		try:
			await self.browser_session.set_extra_headers(CUSTOM_HEADERS, target_id=event.target_id)
		except Exception as e:
			self.logger.debug(f'Could not set headers on {event.target_id[:8]}: {e}')

	async def on_AgentFocusChangedEvent(self, event: AgentFocusChangedEvent) -> None:
		"""Re-apply headers when the agent switches to a different tab."""
		try:
			await self.browser_session.set_extra_headers(CUSTOM_HEADERS, target_id=event.target_id)
		except Exception as e:
			self.logger.debug(f'Could not set headers on {event.target_id[:8]}: {e}')

async def main():
	browser = Browser(headless=False)

	# Start the browser so watchdogs are initialized
	await browser.start()

	# Attach our custom watchdog to the browser session
	CustomHeadersWatchdog.model_rebuild()
	headers_watchdog = CustomHeadersWatchdog(event_bus=browser.event_bus, browser_session=browser)
	headers_watchdog.attach_to_session()

	# The watchdog only fires for tabs created AFTER registration.
	# To apply headers to an already-existing tab, call set_extra_headers():
	#
	#   await browser.set_extra_headers(CUSTOM_HEADERS)
	#   await browser.set_extra_headers(CUSTOM_HEADERS, target_id=some_target_id)
	#
	# Keep in mind that setExtraHTTPHeaders is a full replacement – each
	# call overwrites all previously set extra headers on that target.

	# Run the agent – open httpbin.org/headers in a new tab so the
	# watchdog fires and injects the custom headers.
	agent = Agent(
		task=(
			'Open https://httpbin.org/headers in two different tabs and extract the full JSON response. '
			'Look for the custom headers X-Custom-Auth, X-Request-Source, and X-Trace-Id in the output and compare the results.'
		),
		llm=ChatBrowserUse(model='bu-2-0'),
		browser=browser,
	)

	result = await agent.run()
	print(result.final_result())

	await browser.kill()

if __name__ == '__main__':
	asyncio.run(main())