diff options
Diffstat (limited to 'internal/server')
| -rw-r--r-- | internal/server/handlers/colorwriter.go | 42 | ||||
| -rw-r--r-- | internal/server/handlers/readcommand.go | 26 |
2 files changed, 53 insertions, 15 deletions
diff --git a/internal/server/handlers/colorwriter.go b/internal/server/handlers/colorwriter.go index 6e0f1af..22d36c3 100644 --- a/internal/server/handlers/colorwriter.go +++ b/internal/server/handlers/colorwriter.go @@ -8,15 +8,17 @@ import ( // ColorWriter wraps an io.Writer and applies colors to output type ColorWriter struct { - writer io.Writer - noColor bool + writer io.Writer + noColor bool + isFirstLine bool } // NewColorWriter creates a new ColorWriter func NewColorWriter(writer io.Writer, noColor bool) *ColorWriter { return &ColorWriter{ - writer: writer, - noColor: noColor, + writer: writer, + noColor: noColor, + isFirstLine: true, } } @@ -27,12 +29,26 @@ func (cw *ColorWriter) Write(p []byte) (n int, err error) { return cw.writer.Write(p) } - // Apply colors - coloredStr := brush.Colorfy(string(p)) - coloredBytes := []byte(coloredStr) + // Apply colors to the content + // The brush.Colorfy function will handle the coloring + inputStr := string(p) + coloredStr := brush.Colorfy(inputStr) + + // Add color reset prefix for all lines except the first + var outputBytes []byte + if cw.isFirstLine { + cw.isFirstLine = false + outputBytes = []byte(coloredStr) + } else { + // Add color reset prefix: [39m[49m[49m[39m + colorResetPrefix := "\x1b[39m\x1b[49m\x1b[49m\x1b[39m" + outputBytes = make([]byte, len(colorResetPrefix)+len(coloredStr)) + copy(outputBytes, colorResetPrefix) + copy(outputBytes[len(colorResetPrefix):], coloredStr) + } // Write the colored output - _, err = cw.writer.Write(coloredBytes) + _, err = cw.writer.Write(outputBytes) if err != nil { return 0, err } @@ -40,4 +56,14 @@ func (cw *ColorWriter) Write(p []byte) (n int, err error) { // Return the original byte count to maintain compatibility // (the caller expects n to match len(p)) return len(p), nil +} + +// Close closes the ColorWriter and writes final color reset if needed +func (cw *ColorWriter) Close() error { + if !cw.noColor && config.Client.TermColorsEnable && !cw.isFirstLine { + // Write final color reset line + colorResetPrefix := "\x1b[39m\x1b[49m\x1b[49m\x1b[39m" + cw.writer.Write([]byte(colorResetPrefix)) + } + return nil }
\ No newline at end of file diff --git a/internal/server/handlers/readcommand.go b/internal/server/handlers/readcommand.go index 0a99aaa..901c929 100644 --- a/internal/server/handlers/readcommand.go +++ b/internal/server/handlers/readcommand.go @@ -147,6 +147,7 @@ func (r *readCommand) readFiles(ctx context.Context, ltx lcontext.LContext, var output io.Writer if r.server.serverless { // In serverless mode, write to stdout with color support + // Note: plain mode means no colors, so noColor = plain output = NewColorWriter(os.Stdout, r.server.plain) } else { // In client-server mode, write to server handler lines channel @@ -209,6 +210,11 @@ func (r *readCommand) readFiles(ctx context.Context, ltx lcontext.LContext, } } } + + // Close the output writer if it's a ColorWriter + if colorWriter, ok := output.(*ColorWriter); ok { + colorWriter.Close() + } } // readStdin processes stdin using direct processing @@ -217,6 +223,7 @@ func (r *readCommand) readStdin(ctx context.Context, ltx lcontext.LContext, re r var output io.Writer if r.server.serverless { // In serverless mode, write to stdout with color support + // Note: plain mode means no colors, so noColor = plain output = NewColorWriter(os.Stdout, r.server.plain) } else { // In client-server mode, write to server handler lines channel @@ -236,6 +243,11 @@ func (r *readCommand) readStdin(ctx context.Context, ltx lcontext.LContext, re r r.server.sendln(r.server.serverMessages, dlog.Server.Error(r.server.user, "Error processing stdin", err)) } + + // Close the output writer if it's a ColorWriter + if colorWriter, ok := output.(*ColorWriter); ok { + colorWriter.Close() + } } // isMapReduceCommand checks if this is a command that's part of a MapReduce operation @@ -293,22 +305,22 @@ func (r *readCommand) createProcessor(re regex.Regex, ltx lcontext.LContext, out mapProcessor, err := fs.NewMapProcessor(plain, hostname, queryStr, output) if err != nil { dlog.Server.Error(r.server.user, "Failed to create MapReduce processor", err) - return fs.NewGrepProcessor(re, plain, noColor, hostname, ltx.BeforeContext, ltx.AfterContext, ltx.MaxCount), false + return fs.NewGrepProcessor(re, plain, noColor, hostname, r.server.serverless, ltx.BeforeContext, ltx.AfterContext, ltx.MaxCount), false } return mapProcessor, false } - return fs.NewGrepProcessor(re, plain, noColor, hostname, ltx.BeforeContext, ltx.AfterContext, ltx.MaxCount), false + return fs.NewGrepProcessor(re, plain, noColor, hostname, r.server.serverless, ltx.BeforeContext, ltx.AfterContext, ltx.MaxCount), false case omode.CatClient: if isMapReduce && queryStr != "" { // This is a standalone MapReduce cat operation mapProcessor, err := fs.NewMapProcessor(plain, hostname, queryStr, output) if err != nil { dlog.Server.Error(r.server.user, "Failed to create MapReduce processor", err) - return fs.NewCatProcessor(plain, noColor, hostname), false + return fs.NewCatProcessor(plain, noColor, hostname, r.server.serverless), false } return mapProcessor, false } - return fs.NewCatProcessor(plain, noColor, hostname), false + return fs.NewCatProcessor(plain, noColor, hostname, r.server.serverless), false case omode.TailClient: if isMapReduce && queryStr != "" { // This is a standalone MapReduce tail operation @@ -327,14 +339,14 @@ func (r *readCommand) createProcessor(re regex.Regex, ltx lcontext.LContext, out mapProcessor, err := fs.NewMapProcessor(plain, hostname, queryStr, output) if err != nil { dlog.Server.Error(r.server.user, "Failed to create MapReduce processor", err) - return fs.NewGrepProcessor(re, plain, noColor, hostname, ltx.BeforeContext, ltx.AfterContext, ltx.MaxCount), false + return fs.NewGrepProcessor(re, plain, noColor, hostname, r.server.serverless, ltx.BeforeContext, ltx.AfterContext, ltx.MaxCount), false } return mapProcessor, false } // Fallback - return fs.NewGrepProcessor(re, plain, noColor, hostname, ltx.BeforeContext, ltx.AfterContext, ltx.MaxCount), false + return fs.NewGrepProcessor(re, plain, noColor, hostname, r.server.serverless, ltx.BeforeContext, ltx.AfterContext, ltx.MaxCount), false default: // Default to grep behavior - return fs.NewGrepProcessor(re, plain, noColor, hostname, ltx.BeforeContext, ltx.AfterContext, ltx.MaxCount), false + return fs.NewGrepProcessor(re, plain, noColor, hostname, r.server.serverless, ltx.BeforeContext, ltx.AfterContext, ltx.MaxCount), false } } |