on main help sign in
refactor: simplify context value retrieval using generics

- Introduce generic getContextValue helper function to eliminate code duplication - Reduce code from 75 to 56 lines (25.3% reduction) - Simplify Get*FromContext functions from ~13 lines to 2 lines each - Add comprehensive test coverage (18 test cases) for context functions - Maintain backward compatibility with existing API

commit07d065d780610a4a37e516733da951888293bf4b authorwanghuaiyu@qiniu.com <wanghuaiyu@qiniu.com> dateSun 15 Feb 2026, 22:43 (8 months ago) parent2904a36
2 files changed+236−35
internal/agent/tools/context_test.go+219 internal/agent/tools/tools.go+17−35
internal/agent/tools/context_test.gonew file+219
@@ -0,0 +1,219 @@
1+package tools
2+
3+import (
4+ "context"
5+ "testing"
6+)
7+
8+func TestGetContextValue(t *testing.T) {
9+ tests := []struct {
10+ name string
11+ setup func(ctx context.Context) context.Context
12+ key any
13+ defaultValue any
14+ want any
15+ }{
16+ {
17+ name: "returns string value",
18+ setup: func(ctx context.Context) context.Context {
19+ return context.WithValue(ctx, "testKey", "testValue")
20+ },
21+ key: "testKey",
22+ defaultValue: "",
23+ want: "testValue",
24+ },
25+ {
26+ name: "returns default when key not found",
27+ setup: func(ctx context.Context) context.Context {
28+ return ctx
29+ },
30+ key: "missingKey",
31+ defaultValue: "default",
32+ want: "default",
33+ },
34+ {
35+ name: "returns default when type mismatch",
36+ setup: func(ctx context.Context) context.Context {
37+ return context.WithValue(ctx, "testKey", 123) // int, not string
38+ },
39+ key: "testKey",
40+ defaultValue: "default",
41+ want: "default",
42+ },
43+ {
44+ name: "returns bool value",
45+ setup: func(ctx context.Context) context.Context {
46+ return context.WithValue(ctx, "boolKey", true)
47+ },
48+ key: "boolKey",
49+ defaultValue: false,
50+ want: true,
51+ },
52+ {
53+ name: "returns int value",
54+ setup: func(ctx context.Context) context.Context {
55+ return context.WithValue(ctx, "intKey", 42)
56+ },
57+ key: "intKey",
58+ defaultValue: 0,
59+ want: 42,
60+ },
61+ }
62+
63+ for _, tt := range tests {
64+ t.Run(tt.name, func(t *testing.T) {
65+ ctx := tt.setup(context.Background())
66+
67+ var got any
68+ switch tt.defaultValue.(type) {
69+ case string:
70+ got = getContextValue(ctx, tt.key, tt.defaultValue.(string))
71+ case bool:
72+ got = getContextValue(ctx, tt.key, tt.defaultValue.(bool))
73+ case int:
74+ got = getContextValue(ctx, tt.key, tt.defaultValue.(int))
75+ }
76+
77+ if got != tt.want {
78+ t.Errorf("getContextValue() = %v, want %v", got, tt.want)
79+ }
80+ })
81+ }
82+}
83+
84+func TestGetSessionFromContext(t *testing.T) {
85+ tests := []struct {
86+ name string
87+ ctx context.Context
88+ want string
89+ }{
90+ {
91+ name: "returns session ID when present",
92+ ctx: context.WithValue(context.Background(), SessionIDContextKey, "session-123"),
93+ want: "session-123",
94+ },
95+ {
96+ name: "returns empty string when not present",
97+ ctx: context.Background(),
98+ want: "",
99+ },
100+ {
101+ name: "returns empty string when wrong type",
102+ ctx: context.WithValue(context.Background(), SessionIDContextKey, 123),
103+ want: "",
104+ },
105+ }
106+
107+ for _, tt := range tests {
108+ t.Run(tt.name, func(t *testing.T) {
109+ got := GetSessionFromContext(tt.ctx)
110+ if got != tt.want {
111+ t.Errorf("GetSessionFromContext() = %v, want %v", got, tt.want)
112+ }
113+ })
114+ }
115+}
116+
117+func TestGetMessageFromContext(t *testing.T) {
118+ tests := []struct {
119+ name string
120+ ctx context.Context
121+ want string
122+ }{
123+ {
124+ name: "returns message ID when present",
125+ ctx: context.WithValue(context.Background(), MessageIDContextKey, "msg-456"),
126+ want: "msg-456",
127+ },
128+ {
129+ name: "returns empty string when not present",
130+ ctx: context.Background(),
131+ want: "",
132+ },
133+ {
134+ name: "returns empty string when wrong type",
135+ ctx: context.WithValue(context.Background(), MessageIDContextKey, 456),
136+ want: "",
137+ },
138+ }
139+
140+ for _, tt := range tests {
141+ t.Run(tt.name, func(t *testing.T) {
142+ got := GetMessageFromContext(tt.ctx)
143+ if got != tt.want {
144+ t.Errorf("GetMessageFromContext() = %v, want %v", got, tt.want)
145+ }
146+ })
147+ }
148+}
149+
150+func TestGetSupportsImagesFromContext(t *testing.T) {
151+ tests := []struct {
152+ name string
153+ ctx context.Context
154+ want bool
155+ }{
156+ {
157+ name: "returns true when present and true",
158+ ctx: context.WithValue(context.Background(), SupportsImagesContextKey, true),
159+ want: true,
160+ },
161+ {
162+ name: "returns false when present and false",
163+ ctx: context.WithValue(context.Background(), SupportsImagesContextKey, false),
164+ want: false,
165+ },
166+ {
167+ name: "returns false when not present",
168+ ctx: context.Background(),
169+ want: false,
170+ },
171+ {
172+ name: "returns false when wrong type",
173+ ctx: context.WithValue(context.Background(), SupportsImagesContextKey, "true"),
174+ want: false,
175+ },
176+ }
177+
178+ for _, tt := range tests {
179+ t.Run(tt.name, func(t *testing.T) {
180+ got := GetSupportsImagesFromContext(tt.ctx)
181+ if got != tt.want {
182+ t.Errorf("GetSupportsImagesFromContext() = %v, want %v", got, tt.want)
183+ }
184+ })
185+ }
186+}
187+
188+func TestGetModelNameFromContext(t *testing.T) {
189+ tests := []struct {
190+ name string
191+ ctx context.Context
192+ want string
193+ }{
194+ {
195+ name: "returns model name when present",
196+ ctx: context.WithValue(context.Background(), ModelNameContextKey, "claude-opus-4"),
197+ want: "claude-opus-4",
198+ },
199+ {
200+ name: "returns empty string when not present",
201+ ctx: context.Background(),
202+ want: "",
203+ },
204+ {
205+ name: "returns empty string when wrong type",
206+ ctx: context.WithValue(context.Background(), ModelNameContextKey, 789),
207+ want: "",
208+ },
209+ }
210+
211+ for _, tt := range tests {
212+ t.Run(tt.name, func(t *testing.T) {
213+ got := GetModelNameFromContext(tt.ctx)
214+ if got != tt.want {
215+ t.Errorf("GetModelNameFromContext() = %v, want %v", got, tt.want)
216+ }
217+ })
218+ }
219+}
internal/agent/tools/tools.go+17−35
@@ -22,53 +22,35 @@ const (
2222 ModelNameContextKey modelNameKey = "model_name"
2323)
2424
25−// GetSessionFromContext retrieves the session ID from the context.
26−func GetSessionFromContext(ctx context.Context) string {
27− sessionID := ctx.Value(SessionIDContextKey)
28− if sessionID == nil {
29− return ""
25+// getContextValue is a generic helper that retrieves a typed value from context.
26+// If the value is not found or has the wrong type, it returns the default value.
27+func getContextValue[T any](ctx context.Context, key any, defaultValue T) T {
28+ value := ctx.Value(key)
29+ if value == nil {
30+ return defaultValue
3031 }
31− s, ok := sessionID.(string)
32− if !ok {
33− return ""
32+ if typedValue, ok := value.(T); ok {
33+ return typedValue
3434 }
35− return s
35+ return defaultValue
36+}
37+
38+// GetSessionFromContext retrieves the session ID from the context.
39+func GetSessionFromContext(ctx context.Context) string {
40+ return getContextValue(ctx, SessionIDContextKey, "")
3641}
3742
3843// GetMessageFromContext retrieves the message ID from the context.
3944func GetMessageFromContext(ctx context.Context) string {
40− messageID := ctx.Value(MessageIDContextKey)
41− if messageID == nil {
42− return ""
43− }
44− s, ok := messageID.(string)
45− if !ok {
46− return ""
47− }
48− return s
45+ return getContextValue(ctx, MessageIDContextKey, "")
4946}
5047
5148// GetSupportsImagesFromContext retrieves whether the model supports images from the context.
5249func GetSupportsImagesFromContext(ctx context.Context) bool {
53− supportsImages := ctx.Value(SupportsImagesContextKey)
54− if supportsImages == nil {
55− return false
56− }
57− if supports, ok := supportsImages.(bool); ok {
58− return supports
59− }
60− return false
50+ return getContextValue(ctx, SupportsImagesContextKey, false)
6151}
6252
6353// GetModelNameFromContext retrieves the model name from the context.
6454func GetModelNameFromContext(ctx context.Context) string {
65− modelName := ctx.Value(ModelNameContextKey)
66− if modelName == nil {
67− return ""
68− }
69− s, ok := modelName.(string)
70− if !ok {
71− return ""
72− }
73− return s
55+ return getContextValue(ctx, ModelNameContextKey, "")
7456}