From 8c3fb3270e9e6d9dca1d2272881582515e8c8549 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 21 Apr 2026 23:28:29 +0000 Subject: [PATCH 1/2] fix(model): handle SSE spec leading space and event prefix MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per SSE spec §9.2 the value after "data:" or "event:" may have a single optional leading space that parsers must strip. The previous parser stripped the prefix but not the space, so Gin's "data: [DONE]" was compared as " [DONE]" and the stream never terminated cleanly. Forwarded tokens were also prefixed with stray whitespace. - Extract stripSSEField helper that removes at most one leading space. - Apply it to both event: and data: field names. - Reset isError on blank-line dispatch (event boundary), not inside the data: branch. - Recognize "event:error" in addition to "event: error". https://claude.ai/code/session_01Qa2FSi9HhA49Z4LTMmpxm9 --- model/ai_service.go | 36 ++++++++---- model/ai_service_test.go | 122 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 148 insertions(+), 10 deletions(-) diff --git a/model/ai_service.go b/model/ai_service.go index 31f7964..75f859a 100644 --- a/model/ai_service.go +++ b/model/ai_service.go @@ -72,24 +72,26 @@ func (s *sseAIService) QueryCommandStream( for scanner.Scan() { line := scanner.Text() - if line == "event: error" { - isError = true + if line == "" { + isError = false continue } - if strings.HasPrefix(line, "data:") { - data := line[len("data:"):] + if v, ok := stripSSEField(line, "event:"); ok { + if v == "error" { + isError = true + } + continue + } + if v, ok := stripSSEField(line, "data:"); ok { if isError { - return fmt.Errorf("server error: %s", data) + return fmt.Errorf("server error: %s", v) } - - if data == "[DONE]" { + if v == "[DONE]" { return nil } - - onToken(data) - isError = false + onToken(v) } } @@ -99,3 +101,17 @@ func (s *sseAIService) QueryCommandStream( return nil } + +// stripSSEField returns the value after prefix, stripping one optional leading +// space per the SSE specification (§9.2 "If value starts with a U+0020 SPACE +// character, remove it from value"). +func stripSSEField(line, prefix string) (string, bool) { + if !strings.HasPrefix(line, prefix) { + return "", false + } + v := line[len(prefix):] + if strings.HasPrefix(v, " ") { + v = v[1:] + } + return v, true +} diff --git a/model/ai_service_test.go b/model/ai_service_test.go index 7879557..9e252d0 100644 --- a/model/ai_service_test.go +++ b/model/ai_service_test.go @@ -6,6 +6,7 @@ import ( "fmt" "net/http" "net/http/httptest" + "strings" "testing" ) @@ -90,3 +91,124 @@ func TestQueryCommandStream_ErrorResponseBody(t *testing.T) { }) } } + +func TestQueryCommandStream_SSEParsing(t *testing.T) { + tests := []struct { + name string + body string + wantErr bool + wantErrSubstr string + wantTokens []string + }{ + { + name: "data with space and [DONE] terminates cleanly", + body: "data: [DONE]\n\n", + wantTokens: nil, + }, + { + name: "data without space and [DONE] terminates cleanly", + body: "data:[DONE]\n\n", + wantTokens: nil, + }, + { + name: "single data token with leading space is stripped", + body: "data: hello\n\ndata: [DONE]\n\n", + wantTokens: []string{"hello"}, + }, + { + name: "single data token without leading space passes through", + body: "data:hello\n\ndata:[DONE]\n\n", + wantTokens: []string{"hello"}, + }, + { + name: "multi-token stream concatenates without spurious spaces", + body: "data: ls\n\ndata: -la\n\ndata: [DONE]\n\n", + wantTokens: []string{"ls", " -la"}, + }, + { + name: "event error with space", + body: "event: error\ndata: boom\n\n", + wantErr: true, + wantErrSubstr: "boom", + }, + { + name: "event error without space", + body: "event:error\ndata:boom\n\n", + wantErr: true, + wantErrSubstr: "boom", + }, + { + name: "blank line resets error state between events", + body: "event: error\n\ndata: hello\n\ndata: [DONE]\n\n", + wantTokens: []string{"hello"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(tt.body)) + })) + defer server.Close() + + var got []string + svc := NewAIService() + err := svc.QueryCommandStream( + context.Background(), + CommandSuggestVariables{Shell: "bash", Os: "linux", Query: "test"}, + Endpoint{APIEndpoint: server.URL, Token: "test-token"}, + func(token string) { got = append(got, token) }, + ) + + if tt.wantErr { + if err == nil { + t.Fatalf("expected error, got nil (tokens=%v)", got) + } + if !strings.Contains(err.Error(), tt.wantErrSubstr) { + t.Fatalf("expected error to contain %q, got %q", tt.wantErrSubstr, err.Error()) + } + return + } + + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(got) != len(tt.wantTokens) { + t.Fatalf("token count mismatch: want %d %v, got %d %v", len(tt.wantTokens), tt.wantTokens, len(got), got) + } + for i, tok := range tt.wantTokens { + if got[i] != tok { + t.Errorf("token[%d] = %q, want %q", i, got[i], tok) + } + } + }) + } +} + +func TestStripSSEField(t *testing.T) { + tests := []struct { + name string + line string + prefix string + wantVal string + wantOk bool + }{ + {"no match", "foo:bar", "data:", "", false}, + {"match no space", "data:hello", "data:", "hello", true}, + {"match one space stripped", "data: hello", "data:", "hello", true}, + {"match two spaces preserves second", "data: hello", "data:", " hello", true}, + {"empty value no space", "data:", "data:", "", true}, + {"empty value one space", "data: ", "data:", "", true}, + {"event error with space", "event: error", "event:", "error", true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + v, ok := stripSSEField(tt.line, tt.prefix) + if ok != tt.wantOk || v != tt.wantVal { + t.Errorf("stripSSEField(%q, %q) = (%q, %v), want (%q, %v)", tt.line, tt.prefix, v, ok, tt.wantVal, tt.wantOk) + } + }) + } +} From 5373bf2b76f80e8419e919bbcf6ca647190ce5db Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 21 Apr 2026 23:28:35 +0000 Subject: [PATCH 2/2] fix(cli): sanitize AI command output before execution The suggested command often arrived wrapped in Markdown code fences or surrounding single backticks, which broke shell execution. Extract a sanitizeSuggestedCommand helper that strips triple-backtick fences (with optional bash/sh/shell/zsh/fish/pwsh/powershell language tag) and single-line backtick wrapping, then trims whitespace. Fail fast with a clear user-facing error on empty output instead of classifying an empty string as ActionOther and silently returning. https://claude.ai/code/session_01Qa2FSi9HhA49Z4LTMmpxm9 --- commands/query.go | 39 ++++++++++++++++++++++++++++++++++++++- commands/query_test.go | 37 ++++++++++++++++++++++++++++++++++++- 2 files changed, 74 insertions(+), 2 deletions(-) diff --git a/commands/query.go b/commands/query.go index 7a8f672..8953963 100644 --- a/commands/query.go +++ b/commands/query.go @@ -102,7 +102,11 @@ func commandQuery(c *cli.Context) error { // Print newline after streaming fmt.Println() - newCommand := strings.TrimSpace(result.String()) + newCommand := sanitizeSuggestedCommand(result.String()) + if newCommand == "" { + color.Red.Println("AI returned an empty response. Try rephrasing your query.") + return fmt.Errorf("empty AI response") + } // Check auto-run configuration if cfg.AI != nil && (cfg.AI.Agent.View || cfg.AI.Agent.Edit || cfg.AI.Agent.Delete) { @@ -188,6 +192,39 @@ func executeCommand(ctx context.Context, command string) error { return nil } +// sanitizeSuggestedCommand normalizes raw AI output into an executable command. +// It strips triple-backtick fences (with optional language tag like bash, sh, +// zsh, fish, pwsh, powershell), strips surrounding single backticks when the +// result is a single line, and trims whitespace. Responses that start with `#` +// are treated as refusal comments and preserved verbatim so the caller can +// surface them to the user without attempting execution. +func sanitizeSuggestedCommand(raw string) string { + s := strings.TrimSpace(raw) + if s == "" { + return "" + } + + if strings.HasPrefix(s, "```") { + s = strings.TrimPrefix(s, "```") + if nl := strings.IndexByte(s, '\n'); nl >= 0 { + switch strings.ToLower(strings.TrimSpace(s[:nl])) { + case "", "bash", "sh", "shell", "zsh", "fish", "pwsh", "powershell": + s = s[nl+1:] + } + } + s = strings.TrimRight(s, " \t\n") + s = strings.TrimSuffix(s, "```") + s = strings.TrimSpace(s) + } + + if !strings.ContainsRune(s, '\n') && len(s) >= 2 && + strings.HasPrefix(s, "`") && strings.HasSuffix(s, "`") { + s = strings.TrimSpace(s[1 : len(s)-1]) + } + + return s +} + func getSystemContext(query string) (model.CommandSuggestVariables, error) { // Get shell information shell := os.Getenv("SHELL") diff --git a/commands/query_test.go b/commands/query_test.go index abb13e5..09974b6 100644 --- a/commands/query_test.go +++ b/commands/query_test.go @@ -451,7 +451,42 @@ func (s *queryTestSuite) TestQueryCommandEmptyAIResponse() { } err := s.app.Run(command) - assert.Nil(s.T(), err) + assert.NotNil(s.T(), err) + assert.Contains(s.T(), err.Error(), "empty AI response") +} + +func (s *queryTestSuite) TestSanitizeSuggestedCommand() { + tests := []struct { + name string + in string + want string + }{ + {"plain", "ls -la", "ls -la"}, + {"trims whitespace", " ls -la \n\t", "ls -la"}, + {"fence with bash tag", "```bash\necho hi\n```", "echo hi"}, + {"fence with sh tag", "```sh\necho hi\n```", "echo hi"}, + {"fence with zsh tag", "```zsh\necho hi\n```", "echo hi"}, + {"fence with shell tag", "```shell\necho hi\n```", "echo hi"}, + {"fence with fish tag", "```fish\nset -x FOO bar\n```", "set -x FOO bar"}, + {"fence with powershell tag", "```powershell\nGet-Process\n```", "Get-Process"}, + {"fence with pwsh tag", "```pwsh\nGet-Process\n```", "Get-Process"}, + {"fence no language tag", "```\necho hi\n```", "echo hi"}, + {"fence with trailing newline before closing", "```bash\nls -la\n\n```", "ls -la"}, + {"single backticks around single-line", "`ls -la`", "ls -la"}, + {"single backticks with surrounding space", " `ls -la` ", "ls -la"}, + {"only whitespace", " \n\t ", ""}, + {"empty", "", ""}, + {"comment passthrough preserved", "# refusing: unsafe request", "# refusing: unsafe request"}, + {"multiline without fences kept", "ls\ncat foo", "ls\ncat foo"}, + } + for _, tt := range tests { + s.T().Run(tt.name, func(t *testing.T) { + got := sanitizeSuggestedCommand(tt.in) + if got != tt.want { + t.Errorf("sanitizeSuggestedCommand(%q) = %q, want %q", tt.in, got, tt.want) + } + }) + } } func (s *queryTestSuite) TestQueryCommandDescription() {