Skip to content

Commit b4a8d70

Browse files
committed
Some enhancements
1 parent 3f88185 commit b4a8d70

9 files changed

Lines changed: 174 additions & 23 deletions

File tree

cmd/setup.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ var setupCmd = &cobra.Command{
3939
func init() {
4040
rootCmd.AddCommand(setupCmd)
4141

42-
setupCmd.Flags().StringVar(&setupAPIKey, "api-key", "", "Anthropic API key (or set ANTHROPIC_API_KEY env var) - uses Tusk backend if not provided")
42+
setupCmd.Flags().StringVar(&setupAPIKey, "api-key", "", "Your Anthropic API key (requests go directly to Anthropic). If not provided, uses Tusk's secure proxy")
4343
setupCmd.Flags().StringVar(&setupModel, "model", "claude-sonnet-4-5-20250929", "Claude model to use")
4444
setupCmd.Flags().BoolVar(&setupSkipPermissions, "skip-permissions", false, "Skip permission prompts for consequential actions (commands, file writes, etc.)")
4545
setupCmd.Flags().BoolVar(&setupDisableProgress, "disable-progress-state", false, "Disable progress state (saving to a PROGRESS.md file) or resuming from it")

cmd/short_docs/setup.md

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,9 +14,17 @@ Full docs: <https://docs.usetusk.ai/api-tests/setup-agent>
1414

1515
## Requirements
1616

17-
- Anthropic API key: Set via `ANTHROPIC_API_KEY` environment variable or `--api-key` flag
1817
- Supported languages: Python (FastAPI, Flask, Django, Starlette) and Node.js (Express, Fastify, Koa, Hapi)
1918

19+
## API Key Options
20+
21+
For convenience, by default the setup agent uses Tusk's servers as a proxy to the Anthropic API (no API key needed). Your data is never used for training, see privacy policy: usetusk.ai/privacy.
22+
23+
To use your own Anthropic API key instead:
24+
25+
- Set via `ANTHROPIC_API_KEY` environment variable, or
26+
- Pass `--api-key` flag (requests go directly to Anthropic)
27+
2028
## Modes
2129

2230
- Interactive (default): Shows a TUI with real-time progress and allows you to approve/reject agent actions

internal/agent/agent.go

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -194,7 +194,8 @@ func (a *Agent) Run(parentCtx context.Context) error {
194194
a.ui = NewAgentUI(a.ctx, a.cancel, a.printMode, a.phaseManager.GetPhaseNames(), a.eligibilityOnly)
195195

196196
// Show intro screen and wait for user to continue
197-
shouldContinue, err := a.ui.ShowIntro()
197+
isProxyMode := a.client.mode == APIModeProxy
198+
shouldContinue, err := a.ui.ShowIntro(isProxyMode)
198199
if err != nil {
199200
return fmt.Errorf("failed to show intro: %w", err)
200201
}
@@ -258,6 +259,7 @@ func (a *Agent) runAgent() error {
258259
defer a.cleanup()
259260
a.startTime = time.Now()
260261
a.sessionID = generateSessionID()
262+
a.client.SetSessionID(a.sessionID) // Pass session ID to client for request correlation
261263
a.trackEvent("drift_cli:setup_agent:started", nil)
262264

263265
// Track completed phases for progress file

internal/agent/api.go

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,12 @@ type ClaudeClient struct {
4141
model string
4242
httpClient *http.Client
4343
baseURL string
44+
sessionID string
45+
}
46+
47+
// SetSessionID sets the session ID for request correlation
48+
func (c *ClaudeClient) SetSessionID(sessionID string) {
49+
c.sessionID = sessionID
4450
}
4551

4652
// NewClaudeClient creates a new Claude API client (legacy constructor for BYOK)
@@ -115,7 +121,10 @@ func (c *ClaudeClient) getEndpoint() string {
115121
func (c *ClaudeClient) setAuthHeaders(req *http.Request) {
116122
if c.mode == APIModeProxy {
117123
req.Header.Set("Authorization", "Bearer "+c.bearerToken)
118-
req.Header.Set("X-Tusk-CLI-Version", version.Version)
124+
req.Header.Set("x-tusk-cli-version", version.Version)
125+
if c.sessionID != "" {
126+
req.Header.Set("x-tusk-session-id", c.sessionID)
127+
}
119128
} else {
120129
req.Header.Set("x-api-key", c.apiKey)
121130
req.Header.Set("anthropic-version", "2023-06-01")

internal/agent/intro.go

Lines changed: 30 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ const tuskLogo = `████████╗██╗ ██╗████
2525
██████╔╝██║ ██║██║██║ ██║
2626
╚═════╝ ╚═╝ ╚═╝╚═╝╚═╝ ╚═╝ `
2727

28-
const introDescription = `Welcome to Tusk Drift Setup!
28+
const introDescriptionBase = `Welcome to Tusk Drift Setup!
2929
3030
This AI-powered agent will automatically configure your project
3131
for API testing by analyzing your codebase, installing the Drift
@@ -40,6 +40,18 @@ The agent will guide you through the following phases:
4040
4141
You may be prompted for input during the setup process.`
4242

43+
const introProxyNote = `
44+
45+
Note: For convenience, this setup agent will use Tusk's servers as a proxy to the Anthropic API. Your data is never used for training. See usetusk.ai/privacy for details.`
46+
47+
// getIntroDescription returns the appropriate description based on API mode
48+
func getIntroDescription(isProxyMode bool) string {
49+
if isProxyMode {
50+
return introDescriptionBase + introProxyNote
51+
}
52+
return introDescriptionBase
53+
}
54+
4355
// Color gradient for the wave effect
4456
var waveColors = []string{
4557
"213", // bright magenta
@@ -75,12 +87,15 @@ type IntroModel struct {
7587
logoLines []string
7688
logoWidth int
7789
logoHeight int
90+
91+
// API mode (affects description text)
92+
isProxyMode bool
7893
}
7994

8095
type introTickMsg time.Time
8196

8297
// NewIntroModel creates a new intro screen model
83-
func NewIntroModel() *IntroModel {
98+
func NewIntroModel(isProxyMode bool) *IntroModel {
8499
lines := strings.Split(tuskLogo, "\n")
85100

86101
maxWidth := 0
@@ -92,12 +107,13 @@ func NewIntroModel() *IntroModel {
92107
}
93108

94109
return &IntroModel{
95-
tick: 0,
96-
logoLines: lines,
97-
logoWidth: maxWidth,
98-
logoHeight: len(lines),
99-
mouseX: -1,
100-
mouseY: -1,
110+
tick: 0,
111+
logoLines: lines,
112+
logoWidth: maxWidth,
113+
isProxyMode: isProxyMode,
114+
logoHeight: len(lines),
115+
mouseX: -1,
116+
mouseY: -1,
101117
}
102118
}
103119

@@ -149,7 +165,7 @@ func (m *IntroModel) View() string {
149165
// Content heights
150166
logoHeight := m.logoHeight
151167
spacingAfterLogo := 2
152-
descBoxHeight := 17 // Approx height of description box with borders/padding
168+
descBoxHeight := 20 // Approx height of description box with borders/padding
153169
spacingAfterDesc := 2
154170
footerHeight := 1
155171

@@ -203,7 +219,7 @@ func (m *IntroModel) View() string {
203219
BorderForeground(lipgloss.Color(boxBorderColor)).
204220
Padding(1, 2)
205221

206-
descBox := boxStyle.Render(descStyle.Render(introDescription))
222+
descBox := boxStyle.Render(descStyle.Render(getIntroDescription(m.isProxyMode)))
207223

208224
centeredDesc := lipgloss.NewStyle().
209225
Width(m.width).
@@ -292,8 +308,8 @@ func (m *IntroModel) ShouldContinue() bool {
292308
}
293309

294310
// RunIntroScreen runs the intro screen and returns true if user wants to continue
295-
func RunIntroScreen() (bool, error) {
296-
model := NewIntroModel()
311+
func RunIntroScreen(isProxyMode bool) (bool, error) {
312+
model := NewIntroModel(isProxyMode)
297313
p := tea.NewProgram(model, tea.WithAltScreen(), tea.WithMouseAllMotion())
298314

299315
finalModel, err := p.Run()
@@ -316,7 +332,7 @@ func RunIntroScreen() (bool, error) {
316332
}
317333

318334
// PrintIntroHeadless prints a simple intro for headless mode (no confirmation needed for scripts)
319-
func PrintIntroHeadless() {
335+
func PrintIntroHeadless(isProxyMode bool) {
320336
primaryStyle := lipgloss.NewStyle().
321337
Foreground(lipgloss.Color(styles.PrimaryColor)).
322338
Bold(true)
@@ -327,7 +343,7 @@ func PrintIntroHeadless() {
327343
fmt.Println()
328344
fmt.Println(primaryStyle.Render(tuskLogo))
329345
fmt.Println()
330-
fmt.Println(introDescription)
346+
fmt.Println(getIntroDescription(isProxyMode))
331347
fmt.Println()
332348
fmt.Println(dimStyle.Render("────────────────────────────────────────────────────────"))
333349
fmt.Println()
Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
package agent
2+
3+
import (
4+
"bytes"
5+
"encoding/json"
6+
"io"
7+
"net/http"
8+
"os"
9+
"testing"
10+
"time"
11+
12+
"github.com/Use-Tusk/tusk-drift-cli/internal/api"
13+
"github.com/Use-Tusk/tusk-drift-cli/internal/version"
14+
)
15+
16+
// TestSystemPromptAcceptedByBackend verifies that the current system prompt
17+
// is accepted by the Tusk backend proxy. This test ensures that when the CLI's
18+
// system prompt changes, a corresponding update to the backend's PROMPT_VERSION_RANGES
19+
// has been deployed.
20+
//
21+
// Required environment variables:
22+
// - TUSK_PROXY_TEST_API_KEY: Tusk API key for authentication (does not expire)
23+
//
24+
// Optional environment variables:
25+
// - TUSK_CLOUD_API_URL: Override the backend URL (defaults to https://api.usetusk.ai)
26+
//
27+
// This test uses validate-only mode (x-tusk-validate-only header) to avoid
28+
// Anthropic API calls and associated costs. It only validates that the system
29+
// prompt is accepted by the backend.
30+
//
31+
// This test is skipped if TUSK_PROXY_TEST_API_KEY is not set.
32+
func TestSystemPromptAcceptedByBackend(t *testing.T) {
33+
apiKey := os.Getenv("TUSK_PROXY_TEST_API_KEY")
34+
if apiKey == "" {
35+
t.Skip("Skipping integration test: TUSK_PROXY_TEST_API_KEY must be set")
36+
}
37+
38+
baseURL := api.GetBaseURL() // Defaults to https://api.usetusk.ai
39+
40+
// Build request body using the same struct as production code
41+
requestBody := createMessageRequest{
42+
Model: "claude-sonnet-4-5-20250929",
43+
MaxTokens: 1,
44+
System: SystemPrompt,
45+
Messages: []Message{{Role: "user", Content: []Content{{Type: "text", Text: "test"}}}},
46+
Stream: false,
47+
}
48+
49+
bodyBytes, err := json.Marshal(requestBody)
50+
if err != nil {
51+
t.Fatalf("Failed to marshal request body: %v", err)
52+
}
53+
54+
// Create request
55+
endpoint := baseURL + "/api/drift/setup-agent"
56+
req, err := http.NewRequest("POST", endpoint, bytes.NewReader(bodyBytes))
57+
if err != nil {
58+
t.Fatalf("Failed to create request: %v", err)
59+
}
60+
61+
// Set headers
62+
req.Header.Set("Content-Type", "application/json")
63+
req.Header.Set("x-api-key", apiKey) // API key auth (no expiration)
64+
req.Header.Set("x-tusk-cli-version", version.Version)
65+
req.Header.Set("x-tusk-validate-only", "true") // Skip Anthropic proxy, just validate
66+
67+
// Send request
68+
client := &http.Client{Timeout: 30 * time.Second}
69+
resp, err := client.Do(req)
70+
if err != nil {
71+
t.Fatalf("Request failed: %v", err)
72+
}
73+
defer func() { _ = resp.Body.Close() }()
74+
75+
// Read response body for error details
76+
respBody, _ := io.ReadAll(resp.Body)
77+
78+
// Check for prompt mismatch error (403)
79+
if resp.StatusCode == http.StatusForbidden {
80+
var errorResp struct {
81+
Error string `json:"error"`
82+
Code string `json:"code"`
83+
}
84+
if json.Unmarshal(respBody, &errorResp) == nil && errorResp.Code == "PROMPT_MISMATCH" {
85+
t.Fatalf(`
86+
================================================================================
87+
SYSTEM PROMPT MISMATCH DETECTED
88+
================================================================================
89+
The backend rejected the current system prompt. This means the CLI's system.md
90+
has changed but the backend's PROMPT_VERSION_RANGES has not been updated.
91+
92+
To fix this:
93+
1. Update PROMPT_VERSION_RANGES in tusk/backend/src/api/drift/config/setupAgentPrompts.ts
94+
2. Deploy the backend changes before releasing this CLI version
95+
96+
Error: %s
97+
Code: %s
98+
CLI Version: %s
99+
================================================================================
100+
`, errorResp.Error, errorResp.Code, version.Version)
101+
}
102+
103+
// Other 403 errors (e.g., auth issues)
104+
t.Fatalf("Request returned 403 Forbidden: %s", string(respBody))
105+
}
106+
107+
// Check for other client errors (400-499) that might indicate issues
108+
if resp.StatusCode >= 400 && resp.StatusCode < 500 {
109+
t.Fatalf("Request returned client error %d: %s", resp.StatusCode, string(respBody))
110+
}
111+
112+
// 200 or streaming response means the prompt was accepted
113+
// (We don't need to consume the full response - just knowing it started is enough)
114+
t.Logf("System prompt accepted by backend (status: %d, version: %s)", resp.StatusCode, version.Version)
115+
}

internal/agent/ui.go

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,8 @@ type AgentUI interface {
1111

1212
// ShowIntro displays the intro screen with animation and description.
1313
// Returns true if user wants to continue, false if they quit.
14-
ShowIntro() (bool, error)
14+
// isProxyMode determines whether to show the proxy mode note.
15+
ShowIntro(isProxyMode bool) (bool, error)
1516

1617
// Phase updates
1718
PhaseChange(name, desc string, phaseNum, totalPhases int)

internal/agent/ui_headless.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -59,8 +59,8 @@ func (u *HeadlessUI) Start() error {
5959
}
6060

6161
// ShowIntro displays the intro screen for headless mode (no confirmation for scripts)
62-
func (u *HeadlessUI) ShowIntro() (bool, error) {
63-
PrintIntroHeadless()
62+
func (u *HeadlessUI) ShowIntro(isProxyMode bool) (bool, error) {
63+
PrintIntroHeadless(isProxyMode)
6464
return true, nil
6565
}
6666

internal/agent/ui_tui.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,8 +31,8 @@ func (u *TUIUI) Start() error {
3131
}
3232

3333
// ShowIntro displays the intro screen with animation
34-
func (u *TUIUI) ShowIntro() (bool, error) {
35-
return RunIntroScreen()
34+
func (u *TUIUI) ShowIntro(isProxyMode bool) (bool, error) {
35+
return RunIntroScreen(isProxyMode)
3636
}
3737

3838
// Stop is a no-op for TUI - cleanup happens via the program

0 commit comments

Comments
 (0)