1package fantasy
2
3import (
4 "context"
5 "testing"
6
7 "github.com/stretchr/testify/assert"
8 "github.com/stretchr/testify/require"
9)
10
11func TestAgent_WithUserAgent_PropagatesOnGenerate(t *testing.T) {
12 t.Parallel()
13
14 var capturedCall Call
15 model := &mockLanguageModel{
16 generateFunc: func(_ context.Context, call Call) (*Response, error) {
17 capturedCall = call
18 return &Response{
19 Content: []Content{TextContent{Text: "ok"}},
20 FinishReason: FinishReasonStop,
21 }, nil
22 },
23 }
24
25 agent := NewAgent(model, WithUserAgent("MyApp/2.0"))
26 _, err := agent.Generate(context.Background(), AgentCall{Prompt: "hi"})
27 require.NoError(t, err)
28 assert.Equal(t, "MyApp/2.0", capturedCall.UserAgent)
29 assert.Empty(t, capturedCall.ModelSegment)
30}
31
32func TestAgent_WithModelSegment_PropagatesOnGenerate(t *testing.T) {
33 t.Parallel()
34
35 var capturedCall Call
36 model := &mockLanguageModel{
37 generateFunc: func(_ context.Context, call Call) (*Response, error) {
38 capturedCall = call
39 return &Response{
40 Content: []Content{TextContent{Text: "ok"}},
41 FinishReason: FinishReasonStop,
42 }, nil
43 },
44 }
45
46 agent := NewAgent(model, WithModelSegment("Claude 4.6 Opus"))
47 _, err := agent.Generate(context.Background(), AgentCall{Prompt: "hi"})
48 require.NoError(t, err)
49 assert.Empty(t, capturedCall.UserAgent)
50 assert.Equal(t, "Claude 4.6 Opus", capturedCall.ModelSegment)
51}
52
53func TestAgent_WithUserAgent_PropagatesOnStream(t *testing.T) {
54 t.Parallel()
55
56 var capturedCall Call
57 model := &mockLanguageModel{
58 streamFunc: func(_ context.Context, call Call) (StreamResponse, error) {
59 capturedCall = call
60 return func(yield func(StreamPart) bool) {
61 yield(StreamPart{
62 Type: StreamPartTypeFinish,
63 FinishReason: FinishReasonStop,
64 })
65 }, nil
66 },
67 }
68
69 agent := NewAgent(model, WithUserAgent("StreamApp/1.0"))
70 _, err := agent.Stream(context.Background(), AgentStreamCall{Prompt: "hi"})
71 require.NoError(t, err)
72 assert.Equal(t, "StreamApp/1.0", capturedCall.UserAgent)
73}
74
75func TestAgent_WithModelSegment_PropagatesOnStream(t *testing.T) {
76 t.Parallel()
77
78 var capturedCall Call
79 model := &mockLanguageModel{
80 streamFunc: func(_ context.Context, call Call) (StreamResponse, error) {
81 capturedCall = call
82 return func(yield func(StreamPart) bool) {
83 yield(StreamPart{
84 Type: StreamPartTypeFinish,
85 FinishReason: FinishReasonStop,
86 })
87 }, nil
88 },
89 }
90
91 agent := NewAgent(model, WithModelSegment("GPT-5"))
92 _, err := agent.Stream(context.Background(), AgentStreamCall{Prompt: "hi"})
93 require.NoError(t, err)
94 assert.Equal(t, "GPT-5", capturedCall.ModelSegment)
95}
96
97func TestAgent_NoUA_OmitsCallLevelFields(t *testing.T) {
98 t.Parallel()
99
100 var capturedCall Call
101 model := &mockLanguageModel{
102 generateFunc: func(_ context.Context, call Call) (*Response, error) {
103 capturedCall = call
104 return &Response{
105 Content: []Content{TextContent{Text: "ok"}},
106 FinishReason: FinishReasonStop,
107 }, nil
108 },
109 }
110
111 agent := NewAgent(model)
112 _, err := agent.Generate(context.Background(), AgentCall{Prompt: "hi"})
113 require.NoError(t, err)
114 assert.Empty(t, capturedCall.UserAgent)
115 assert.Empty(t, capturedCall.ModelSegment)
116}
117
118func TestAgent_WithUserAgentAndModelSegment_BothPropagate(t *testing.T) {
119 t.Parallel()
120
121 var capturedCall Call
122 model := &mockLanguageModel{
123 generateFunc: func(_ context.Context, call Call) (*Response, error) {
124 capturedCall = call
125 return &Response{
126 Content: []Content{TextContent{Text: "ok"}},
127 FinishReason: FinishReasonStop,
128 }, nil
129 },
130 }
131
132 agent := NewAgent(model, WithUserAgent("App/1.0"), WithModelSegment("Claude 4.6"))
133 _, err := agent.Generate(context.Background(), AgentCall{Prompt: "hi"})
134 require.NoError(t, err)
135 assert.Equal(t, "App/1.0", capturedCall.UserAgent)
136 assert.Equal(t, "Claude 4.6", capturedCall.ModelSegment)
137}