Fix ANSI color preservation across multiple log lines

Adds ANSI state tracking to the presenter struct,
restoring color codes at the start of each line and properly
resetting them at line end.

Signed-off-by: neo <[email protected]>

Author: chojs23

cmd/formatter/logs.go

@@ -20,6 +20,7 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"regexp"
 	"strconv"
 	"strings"
 	"sync"
@@ -119,13 +120,73 @@ func (l *logConsumer) write(w io.Writer, container, message string) {
 	}
 	p := l.getPresenter(container)
 	timestamp := time.Now().Format(jsonmessage.RFC3339NanoFixed)
-	for _, line := range strings.Split(message, "\n") {
+
+	lines := strings.Split(message, "\n")
+
+	for _, line := range lines {
+		formattedLine := line
+		if p.ansiState != "" {
+			formattedLine = p.ansiState + line
+		}
+
 		if l.timestamp {
-			_, _ = fmt.Fprintf(w, "%s%s %s\n", p.prefix, timestamp, line)
+			_, _ = fmt.Fprintf(w, "%s%s %s", p.prefix, timestamp, formattedLine)
 		} else {
-			_, _ = fmt.Fprintf(w, "%s%s\n", p.prefix, line)
+			_, _ = fmt.Fprintf(w, "%s%s", p.prefix, formattedLine)
 		}
+
+		if p.ansiState != "" || hasANSICodes(line) {
+			_, _ = fmt.Fprint(w, "\033[0m")
+		}
+		_, _ = fmt.Fprint(w, "\n")
+
+		p.ansiState = extractANSIState(formattedLine)
+	}
+}
+
+var ansiSGRPattern = regexp.MustCompile(`\033\[([0-9;]*)m`)
+
+func hasANSICodes(s string) bool {
+	return ansiSGRPattern.MatchString(s)
+}
+
+func extractANSIState(line string) string {
+	matches := ansiSGRPattern.FindAllStringSubmatch(line, -1)
+	if len(matches) == 0 {
+		return ""
 	}
+
+	state := make(map[string]bool)
+	var activeFormats []string
+
+	for _, match := range matches {
+		codes := match[1]
+		if codes == "" || codes == "0" {
+			state = make(map[string]bool)
+			activeFormats = nil
+			continue
+		}
+
+		parts := strings.Split(codes, ";")
+		for _, part := range parts {
+			if part == "0" {
+				state = make(map[string]bool)
+				activeFormats = nil
+			} else {
+				state[part] = true
+			}
+		}
+	}
+
+	for code := range state {
+		activeFormats = append(activeFormats, code)
+	}
+
+	if len(activeFormats) == 0 {
+		return ""
+	}
+
+	return fmt.Sprintf("\033[%sm", strings.Join(activeFormats, ";"))
 }
 
 func (l *logConsumer) Status(container, msg string) {
@@ -147,9 +208,10 @@ func (l *logConsumer) computeWidth() {
 }
 
 type presenter struct {
-	colors colorFunc
-	name   string
-	prefix string
+	colors    colorFunc
+	name      string
+	prefix    string
+	ansiState string
 }
 
 func (p *presenter) setPrefix(width int) {

cmd/formatter/logs_test.go

@@ -0,0 +1,177 @@
+/*
+   Copyright 2020 Docker Compose CLI authors
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+package formatter
+
+import (
+	"bytes"
+	"context"
+	"strings"
+	"testing"
+
+	"gotest.tools/v3/assert"
+)
+
+func TestANSIStatePreservation(t *testing.T) {
+	tests := []struct {
+		name     string
+		input    string
+		expected []string
+	}{
+		{
+			name:  "red color across multiple lines",
+			input: "\033[31mThis line is RED.\nThis line is also RED.\033[0m",
+			expected: []string{
+				"This line is RED.",
+				"This line is also RED.",
+			},
+		},
+		{
+			name:  "color change within multiline",
+			input: "\033[31mThis is RED.\nStill RED.\nNow \033[34mBLUE.\033[0m",
+			expected: []string{
+				"This is RED.",
+				"Still RED.",
+				"Now \033[34mBLUE.",
+			},
+		},
+		{
+			name:  "no ANSI codes",
+			input: "Plain text\nMore plain text",
+			expected: []string{
+				"Plain text",
+				"More plain text",
+			},
+		},
+		{
+			name:  "single line with ANSI",
+			input: "\033[32mGreen text\033[0m",
+			expected: []string{
+				"Green text",
+			},
+		},
+		{
+			name:  "reset in middle of multiline",
+			input: "\033[31mRed\nStill red\033[0m\nNow normal\nStill normal",
+			expected: []string{
+				"Red",
+				"Still red",
+				"Now normal",
+				"Still normal",
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			buf := &bytes.Buffer{}
+			consumer := NewLogConsumer(context.Background(), buf, buf, false, false, false)
+			consumer.Log("test", tt.input)
+
+			output := buf.String()
+			lines := strings.Split(strings.TrimSuffix(output, "\n"), "\n")
+
+			assert.Equal(t, len(tt.expected), len(lines), "number of lines should match")
+
+			for i, expectedContent := range tt.expected {
+				lineWithoutANSI := stripANSIExceptContent(lines[i])
+				assert.Assert(t, strings.Contains(lineWithoutANSI, expectedContent),
+					"line %d should contain expected content. got: %q, want to contain: %q",
+					i, lineWithoutANSI, expectedContent)
+			}
+		})
+	}
+}
+
+func TestExtractANSIState(t *testing.T) {
+	tests := []struct {
+		name     string
+		input    string
+		expected string
+	}{
+		{
+			name:     "red color code",
+			input:    "\033[31mRed text",
+			expected: "\033[31m",
+		},
+		{
+			name:     "reset code",
+			input:    "\033[31mRed\033[0m",
+			expected: "",
+		},
+		{
+			name:     "no ANSI codes",
+			input:    "Plain text",
+			expected: "",
+		},
+		{
+			name:     "multiple codes",
+			input:    "\033[1m\033[31mBold red",
+			expected: "\033[1;31m",
+		},
+		{
+			name:     "code then reset",
+			input:    "\033[31mRed\033[0mNormal",
+			expected: "",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := extractANSIState(tt.input)
+			if tt.expected == "" {
+				assert.Equal(t, "", result)
+			} else {
+				assert.Assert(t, result != "", "expected non-empty ANSI state")
+			}
+		})
+	}
+}
+
+func TestHasANSICodes(t *testing.T) {
+	tests := []struct {
+		name     string
+		input    string
+		expected bool
+	}{
+		{
+			name:     "with ANSI codes",
+			input:    "\033[31mRed text\033[0m",
+			expected: true,
+		},
+		{
+			name:     "no ANSI codes",
+			input:    "Plain text",
+			expected: false,
+		},
+		{
+			name:     "empty string",
+			input:    "",
+			expected: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := hasANSICodes(tt.input)
+			assert.Equal(t, tt.expected, result)
+		})
+	}
+}
+
+func stripANSIExceptContent(s string) string {
+	return strings.TrimSpace(s)
+}