Skip to content

Commit f1ad7ea

Browse files
Claudelpcox
Claude
and
lpcox
committed
Refactor: consolidate function clustering - move Docker utils, replace stdlib reimplementation, add strutil package
Co-authored-by: lpcox <15877973+lpcox@users.noreply.github.com>
1 parent 9cf3404 commit f1ad7ea

7 files changed

Lines changed: 272 additions & 132 deletions

File tree

‎internal/dockerutil/env.go‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
package dockerutil
2+
3+
import (
4+
"fmt"
5+
"os"
6+
"strings"
7+
)
8+
9+
// ExpandEnvArgs expands Docker -e flags that reference environment variables
10+
// Converts "-e VAR_NAME" to "-e VAR_NAME=value" by reading from the process environment
11+
func ExpandEnvArgs(args []string) []string {
12+
result := make([]string, 0, len(args))
13+
for i := 0; i < len(args); i++ {
14+
arg := args[i]
15+
16+
// Check if this is a -e flag
17+
if arg == "-e" && i+1 < len(args) {
18+
nextArg := args[i+1]
19+
// If next arg doesn't contain '=', it's a variable reference
20+
if len(nextArg) > 0 && !strings.Contains(nextArg, "=") {
21+
// Look up the variable in the environment
22+
if value, exists := os.LookupEnv(nextArg); exists {
23+
result = append(result, "-e")
24+
result = append(result, fmt.Sprintf("%s=%s", nextArg, value))
25+
i++ // Skip the next arg since we processed it
26+
continue
27+
}
28+
}
29+
}
30+
result = append(result, arg)
31+
}
32+
return result
33+
}

‎internal/dockerutil/env_test.go‎

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
package dockerutil
2+
3+
import (
4+
"os"
5+
"testing"
6+
)
7+
8+
func TestExpandEnvArgs(t *testing.T) {
9+
tests := []struct {
10+
name string
11+
args []string
12+
envVars map[string]string
13+
expected []string
14+
}{
15+
{
16+
name: "no -e flags",
17+
args: []string{"run", "--rm", "image"},
18+
envVars: map[string]string{},
19+
expected: []string{"run", "--rm", "image"},
20+
},
21+
{
22+
name: "expand single env variable",
23+
args: []string{"run", "-e", "VAR_NAME", "image"},
24+
envVars: map[string]string{"VAR_NAME": "value1"},
25+
expected: []string{"run", "-e", "VAR_NAME=value1", "image"},
26+
},
27+
{
28+
name: "expand multiple env variables",
29+
args: []string{"run", "-e", "VAR1", "-e", "VAR2", "image"},
30+
envVars: map[string]string{"VAR1": "value1", "VAR2": "value2"},
31+
expected: []string{"run", "-e", "VAR1=value1", "-e", "VAR2=value2", "image"},
32+
},
33+
{
34+
name: "preserve existing key=value format",
35+
args: []string{"run", "-e", "VAR=predefined", "image"},
36+
envVars: map[string]string{},
37+
expected: []string{"run", "-e", "VAR=predefined", "image"},
38+
},
39+
{
40+
name: "mixed: expand and preserve",
41+
args: []string{"run", "-e", "VAR1", "-e", "VAR2=fixed", "image"},
42+
envVars: map[string]string{"VAR1": "value1"},
43+
expected: []string{"run", "-e", "VAR1=value1", "-e", "VAR2=fixed", "image"},
44+
},
45+
{
46+
name: "undefined env variable",
47+
args: []string{"run", "-e", "UNDEFINED_VAR", "image"},
48+
envVars: map[string]string{},
49+
expected: []string{"run", "-e", "UNDEFINED_VAR", "image"},
50+
},
51+
{
52+
name: "empty env variable value",
53+
args: []string{"run", "-e", "EMPTY_VAR", "image"},
54+
envVars: map[string]string{"EMPTY_VAR": ""},
55+
expected: []string{"run", "-e", "EMPTY_VAR=", "image"},
56+
},
57+
{
58+
name: "-e at end of args (no following arg)",
59+
args: []string{"run", "image", "-e"},
60+
envVars: map[string]string{},
61+
expected: []string{"run", "image", "-e"},
62+
},
63+
}
64+
65+
for _, tt := range tests {
66+
t.Run(tt.name, func(t *testing.T) {
67+
// Set up environment variables for test
68+
for k, v := range tt.envVars {
69+
os.Setenv(k, v)
70+
}
71+
// Clean up after test
72+
t.Cleanup(func() {
73+
for k := range tt.envVars {
74+
os.Unsetenv(k)
75+
}
76+
})
77+
78+
result := ExpandEnvArgs(tt.args)
79+
80+
if len(result) != len(tt.expected) {
81+
t.Fatalf("Expected %d args, got %d: %v", len(tt.expected), len(result), result)
82+
}
83+
84+
for i := range result {
85+
if result[i] != tt.expected[i] {
86+
t.Errorf("Arg %d: expected '%s', got '%s'", i, tt.expected[i], result[i])
87+
}
88+
}
89+
})
90+
}
91+
}

‎internal/logger/rpc_helpers.go‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import (
2020
"strings"
2121

2222
"github.com/github/gh-aw-mcpg/internal/logger/sanitize"
23+
"github.com/github/gh-aw-mcpg/internal/strutil"
2324
)
2425

2526
// Pre-compiled regexes for performance (avoid recompiling in hot paths).
@@ -44,10 +45,7 @@ func truncateAndSanitize(payload string, maxLength int) string {
4445
sanitized := sanitize.SanitizeString(payload)
4546

4647
// Then truncate if needed
47-
if len(sanitized) > maxLength {
48-
return sanitized[:maxLength] + "..."
49-
}
50-
return sanitized
48+
return strutil.Truncate(sanitized, maxLength)
5149
}
5250

5351
// extractEssentialFields extracts key fields from the payload for logging

‎internal/mcp/connection.go‎

Lines changed: 4 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -10,14 +10,15 @@ import (
1010
"log"
1111
"log/slog"
1212
"net/http"
13-
"os"
1413
"os/exec"
1514
"strings"
1615
"sync/atomic"
1716
"time"
1817

18+
"github.com/github/gh-aw-mcpg/internal/dockerutil"
1919
"github.com/github/gh-aw-mcpg/internal/logger"
2020
"github.com/github/gh-aw-mcpg/internal/logger/sanitize"
21+
"github.com/github/gh-aw-mcpg/internal/strutil"
2122
"github.com/github/gh-aw-mcpg/internal/version"
2223
sdk "github.com/modelcontextprotocol/go-sdk/mcp"
2324
)
@@ -93,10 +94,7 @@ func parseJSONRPCResponseWithSSE(body []byte, statusCode int, contextDesc string
9394
return httpErrorResponse(), nil
9495
}
9596
// Include the response body to help debug what the server actually returned
96-
bodyPreview := string(body)
97-
if len(bodyPreview) > 500 {
98-
bodyPreview = bodyPreview[:500] + "... (truncated)"
99-
}
97+
bodyPreview := strutil.TruncateWithSuffix(string(body), 500, "... (truncated)")
10098
return nil, fmt.Errorf("failed to parse %s (received non-JSON or malformed response): %w\nResponse body: %s", contextDesc, sseErr, bodyPreview)
10199
}
102100

@@ -314,7 +312,7 @@ func NewConnection(ctx context.Context, serverID, command string, args []string,
314312

315313
// Expand Docker -e flags that reference environment variables
316314
// Docker's `-e VAR_NAME` expects VAR_NAME to be in the environment
317-
expandedArgs := expandDockerEnvArgs(args)
315+
expandedArgs := dockerutil.ExpandEnvArgs(args)
318316
logConn.Printf("Expanded args for Docker env: %v", sanitize.SanitizeArgs(expandedArgs))
319317

320318
// Create command transport
@@ -965,41 +963,6 @@ func (c *Connection) getPrompt(params interface{}) (*Response, error) {
965963
return marshalToResponse(result)
966964
}
967965

968-
// expandDockerEnvArgs expands Docker -e flags that reference environment variables
969-
// Converts "-e VAR_NAME" to "-e VAR_NAME=value" by reading from the process environment
970-
func expandDockerEnvArgs(args []string) []string {
971-
result := make([]string, 0, len(args))
972-
for i := 0; i < len(args); i++ {
973-
arg := args[i]
974-
975-
// Check if this is a -e flag
976-
if arg == "-e" && i+1 < len(args) {
977-
nextArg := args[i+1]
978-
// If next arg doesn't contain '=', it's a variable reference
979-
if len(nextArg) > 0 && !containsEqual(nextArg) {
980-
// Look up the variable in the environment
981-
if value, exists := os.LookupEnv(nextArg); exists {
982-
result = append(result, "-e")
983-
result = append(result, fmt.Sprintf("%s=%s", nextArg, value))
984-
i++ // Skip the next arg since we processed it
985-
continue
986-
}
987-
}
988-
}
989-
result = append(result, arg)
990-
}
991-
return result
992-
}
993-
994-
func containsEqual(s string) bool {
995-
for _, c := range s {
996-
if c == '=' {
997-
return true
998-
}
999-
}
1000-
return false
1001-
}
1002-
1003966
// Close closes the connection
1004967
func (c *Connection) Close() error {
1005968
c.cancel()

‎internal/mcp/connection_test.go‎

Lines changed: 0 additions & 87 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@ import (
77
"io"
88
"net/http"
99
"net/http/httptest"
10-
"os"
1110
"strings"
1211
"testing"
1312

@@ -131,92 +130,6 @@ func TestHTTPRequest_ConfiguredHeaders(t *testing.T) {
131130
assert.Equal(t, sessionID, receivedSessionID)
132131
}
133132

134-
// TestExpandDockerEnvArgs tests the Docker environment variable expansion function
135-
func TestExpandDockerEnvArgs(t *testing.T) {
136-
tests := []struct {
137-
name string
138-
args []string
139-
envVars map[string]string
140-
expected []string
141-
}{
142-
{
143-
name: "no -e flags",
144-
args: []string{"run", "--rm", "image"},
145-
envVars: map[string]string{},
146-
expected: []string{"run", "--rm", "image"},
147-
},
148-
{
149-
name: "expand single env variable",
150-
args: []string{"run", "-e", "VAR_NAME", "image"},
151-
envVars: map[string]string{"VAR_NAME": "value1"},
152-
expected: []string{"run", "-e", "VAR_NAME=value1", "image"},
153-
},
154-
{
155-
name: "expand multiple env variables",
156-
args: []string{"run", "-e", "VAR1", "-e", "VAR2", "image"},
157-
envVars: map[string]string{"VAR1": "value1", "VAR2": "value2"},
158-
expected: []string{"run", "-e", "VAR1=value1", "-e", "VAR2=value2", "image"},
159-
},
160-
{
161-
name: "preserve existing key=value format",
162-
args: []string{"run", "-e", "VAR=predefined", "image"},
163-
envVars: map[string]string{},
164-
expected: []string{"run", "-e", "VAR=predefined", "image"},
165-
},
166-
{
167-
name: "mixed: expand and preserve",
168-
args: []string{"run", "-e", "VAR1", "-e", "VAR2=fixed", "image"},
169-
envVars: map[string]string{"VAR1": "value1"},
170-
expected: []string{"run", "-e", "VAR1=value1", "-e", "VAR2=fixed", "image"},
171-
},
172-
{
173-
name: "undefined env variable",
174-
args: []string{"run", "-e", "UNDEFINED_VAR", "image"},
175-
envVars: map[string]string{},
176-
expected: []string{"run", "-e", "UNDEFINED_VAR", "image"},
177-
},
178-
{
179-
name: "empty env variable value",
180-
args: []string{"run", "-e", "EMPTY_VAR", "image"},
181-
envVars: map[string]string{"EMPTY_VAR": ""},
182-
expected: []string{"run", "-e", "EMPTY_VAR=", "image"},
183-
},
184-
{
185-
name: "-e at end of args (no following arg)",
186-
args: []string{"run", "image", "-e"},
187-
envVars: map[string]string{},
188-
expected: []string{"run", "image", "-e"},
189-
},
190-
}
191-
192-
for _, tt := range tests {
193-
t.Run(tt.name, func(t *testing.T) {
194-
// Set up environment variables for test
195-
for k, v := range tt.envVars {
196-
os.Setenv(k, v)
197-
}
198-
// Clean up after test
199-
t.Cleanup(func() {
200-
for k := range tt.envVars {
201-
os.Unsetenv(k)
202-
}
203-
})
204-
205-
result := expandDockerEnvArgs(tt.args)
206-
207-
if len(result) != len(tt.expected) {
208-
t.Fatalf("Expected %d args, got %d: %v", len(tt.expected), len(result), result)
209-
}
210-
211-
for i := range result {
212-
if result[i] != tt.expected[i] {
213-
t.Errorf("Arg %d: expected '%s', got '%s'", i, tt.expected[i], result[i])
214-
}
215-
}
216-
})
217-
}
218-
}
219-
220133
// TestHTTPRequest_ErrorResponses tests handling of various error conditions
221134
func TestHTTPRequest_ErrorResponses(t *testing.T) {
222135
tests := []struct {

‎internal/strutil/truncate.go‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
package strutil
2+
3+
// Truncate truncates a string to the specified maximum length.
4+
// If the string is longer than maxLen, it's truncated and "..." is appended.
5+
// If maxLen is 0, returns "..." for non-empty strings, empty string for empty strings.
6+
// If maxLen is negative, the original string is returned.
7+
func Truncate(s string, maxLen int) string {
8+
if maxLen < 0 {
9+
return s
10+
}
11+
if maxLen == 0 {
12+
if len(s) > 0 {
13+
return "..."
14+
}
15+
return ""
16+
}
17+
if len(s) <= maxLen {
18+
return s
19+
}
20+
return s[:maxLen] + "..."
21+
}
22+
23+
// TruncateWithSuffix truncates a string to the specified maximum length with a custom suffix.
24+
// If the string is longer than maxLen, it's truncated and suffix is appended.
25+
// If maxLen is 0 or negative, the original string is returned.
26+
func TruncateWithSuffix(s string, maxLen int, suffix string) string {
27+
if maxLen <= 0 || len(s) <= maxLen {
28+
return s
29+
}
30+
return s[:maxLen] + suffix
31+
}

0 commit comments

Comments
 (0)