Merge pull request #907 from rumpl/shell-render

Tool rendering niceties

Author: dgageot

pkg/tui/components/tool/defaulttool/defaulttool.go

@@ -1,10 +1,7 @@
 package defaulttool
 
 import (
-	"fmt"
-
 	tea "charm.land/bubbletea/v2"
-	"github.com/charmbracelet/glamour/v2"
 
 	"github.com/docker/cagent/pkg/tui/components/spinner"
 	"github.com/docker/cagent/pkg/tui/components/toolcommon"
@@ -18,25 +15,22 @@ import (
 // that don't have a specialized component registered.
 // It provides a standard visualization with tool name, arguments, and results.
 type Component struct {
-	message  *types.Message
-	renderer *glamour.TermRenderer
-	spinner  spinner.Spinner
-	width    int
-	height   int
+	message *types.Message
+	spinner spinner.Spinner
+	width   int
+	height  int
 }
 
 // New creates a new default tool component.
 func New(
 	msg *types.Message,
-	renderer *glamour.TermRenderer,
 	_ *service.SessionState,
 ) layout.Model {
 	return &Component{
-		message:  msg,
-		renderer: renderer,
-		spinner:  spinner.New(spinner.ModeSpinnerOnly),
-		width:    80,
-		height:   1,
+		message: msg,
+		spinner: spinner.New(spinner.ModeSpinnerOnly),
+		width:   80,
+		height:  1,
 	}
 }
 
@@ -68,20 +62,20 @@ func (c *Component) Update(msg tea.Msg) (layout.Model, tea.Cmd) {
 func (c *Component) View() string {
 	msg := c.message
 	displayName := msg.ToolDefinition.DisplayName()
-	content := fmt.Sprintf("%s %s", toolcommon.Icon(msg.ToolStatus), styles.HighlightStyle.Render(displayName))
 
-	if msg.ToolStatus == types.ToolStatusPending || msg.ToolStatus == types.ToolStatusRunning {
-		content += " " + c.spinner.View()
+	var argsContent string
+	if msg.ToolCall.Function.Arguments != "" {
+		argsContent = renderToolArgs(msg.ToolCall, c.width-3)
 	}
 
-	if msg.ToolCall.Function.Arguments != "" {
-		content += "\n" + renderToolArgs(msg.ToolCall, c.width-3)
+	if argsContent == "" {
+		return toolcommon.RenderTool(toolcommon.Icon(msg.ToolStatus), msg.ToolDefinition.DisplayName(), c.spinner.View(), "", c.width)
 	}
 
 	var resultContent string
 	if (msg.ToolStatus == types.ToolStatusCompleted || msg.ToolStatus == types.ToolStatusError) && msg.Content != "" {
 		resultContent = toolcommon.FormatToolResult(msg.Content, c.width)
 	}
 
-	return styles.BaseStyle.PaddingLeft(2).PaddingTop(1).Render(content + resultContent)
+	return toolcommon.RenderTool(toolcommon.Icon(msg.ToolStatus), styles.HighlightStyle.Render(displayName), argsContent, resultContent, c.width)
 }

pkg/tui/components/tool/editfile/editfile.go

@@ -5,7 +5,6 @@ import (
 	"fmt"
 
 	tea "charm.land/bubbletea/v2"
-	"github.com/charmbracelet/glamour/v2"
 
 	"github.com/docker/cagent/pkg/tools/builtin"
 	"github.com/docker/cagent/pkg/tui/components/spinner"
@@ -21,7 +20,6 @@ type ToggleDiffViewMsg struct{}
 // Component is a specialized component for rendering edit_file tool calls.
 type Component struct {
 	message      *types.Message
-	renderer     *glamour.TermRenderer
 	spinner      spinner.Spinner
 	width        int
 	height       int
@@ -30,12 +28,10 @@ type Component struct {
 
 func New(
 	msg *types.Message,
-	renderer *glamour.TermRenderer,
 	sessionState *service.SessionState,
 ) layout.Model {
 	return &Component{
 		message:      msg,
-		renderer:     renderer,
 		spinner:      spinner.New(spinner.ModeSpinnerOnly),
 		width:        80,
 		height:       1,
@@ -91,5 +87,5 @@ func (c *Component) View() string {
 		resultContent = toolcommon.FormatToolResult(msg.Content, c.width)
 	}
 
-	return styles.BaseStyle.PaddingLeft(2).PaddingTop(1).Render(content + resultContent)
+	return styles.BaseStyle.Render(content + resultContent)
 }

pkg/tui/components/tool/factory.go

@@ -7,6 +7,7 @@ import (
 	"github.com/docker/cagent/pkg/tui/components/tool/defaulttool"
 	"github.com/docker/cagent/pkg/tui/components/tool/editfile"
 	"github.com/docker/cagent/pkg/tui/components/tool/readfile"
+	"github.com/docker/cagent/pkg/tui/components/tool/shell"
 	"github.com/docker/cagent/pkg/tui/components/tool/todotool"
 	"github.com/docker/cagent/pkg/tui/components/tool/transfertask"
 	"github.com/docker/cagent/pkg/tui/components/tool/writefile"
@@ -36,10 +37,10 @@ func (f *Factory) Create(
 	toolName := msg.ToolCall.Function.Name
 
 	if builder, ok := f.registry.Get(toolName); ok {
-		return builder(msg, renderer, sessionState)
+		return builder(msg, sessionState)
 	}
 
-	return defaulttool.New(msg, renderer, sessionState)
+	return defaulttool.New(msg, sessionState)
 }
 
 var (
@@ -58,6 +59,7 @@ func newDefaultRegistry() *Registry {
 	registry.Register(builtin.ToolNameCreateTodos, todotool.New)
 	registry.Register(builtin.ToolNameUpdateTodo, todotool.New)
 	registry.Register(builtin.ToolNameListTodos, todotool.New)
+	registry.Register(builtin.ToolNameShell, shell.New)
 
 	return registry
 }

pkg/tui/components/tool/readfile/readfile.go

@@ -2,10 +2,8 @@ package readfile
 
 import (
 	"encoding/json"
-	"fmt"
 
 	tea "charm.land/bubbletea/v2"
-	"github.com/charmbracelet/glamour/v2"
 
 	"github.com/docker/cagent/pkg/tools/builtin"
 	"github.com/docker/cagent/pkg/tui/components/spinner"
@@ -27,7 +25,6 @@ type Component struct {
 // New creates a new read file component.
 func New(
 	msg *types.Message,
-	_ *glamour.TermRenderer,
 	_ *service.SessionState,
 ) layout.Model {
 	return &Component{
@@ -67,15 +64,8 @@ func (c *Component) View() string {
 	msg := c.message
 	var args builtin.ReadFileArgs
 	if err := json.Unmarshal([]byte(msg.ToolCall.Function.Arguments), &args); err != nil {
-		return ""
+		return toolcommon.RenderTool(toolcommon.Icon(msg.ToolStatus), msg.ToolDefinition.DisplayName(), c.spinner.View(), "", c.width)
 	}
 
-	displayName := msg.ToolDefinition.DisplayName()
-	content := fmt.Sprintf("%s %s %s", toolcommon.Icon(msg.ToolStatus), styles.HighlightStyle.Render(displayName), styles.MutedStyle.Render(args.Path))
-
-	if msg.ToolStatus == types.ToolStatusPending || msg.ToolStatus == types.ToolStatusRunning {
-		content += " " + c.spinner.View()
-	}
-
-	return styles.BaseStyle.PaddingLeft(2).Render(content)
+	return toolcommon.RenderTool(toolcommon.Icon(msg.ToolStatus), msg.ToolDefinition.DisplayName(), styles.MutedStyle.Render(args.Path), "", c.width)
 }

pkg/tui/components/tool/registry.go

@@ -3,8 +3,6 @@ package tool
 import (
 	"sync"
 
-	"github.com/charmbracelet/glamour/v2"
-
 	"github.com/docker/cagent/pkg/tui/core/layout"
 	"github.com/docker/cagent/pkg/tui/service"
 	"github.com/docker/cagent/pkg/tui/types"
@@ -13,7 +11,6 @@ import (
 // ComponentBuilder is a function that creates a tool component.
 type ComponentBuilder func(
 	msg *types.Message,
-	renderer *glamour.TermRenderer,
 	sessionState *service.SessionState,
 ) layout.Model
 

pkg/tui/components/tool/shell/shell.go

@@ -0,0 +1,69 @@
+package shell
+
+import (
+	"encoding/json"
+
+	tea "charm.land/bubbletea/v2"
+
+	"github.com/docker/cagent/pkg/tools/builtin"
+	"github.com/docker/cagent/pkg/tui/components/spinner"
+	"github.com/docker/cagent/pkg/tui/components/toolcommon"
+	"github.com/docker/cagent/pkg/tui/core/layout"
+	"github.com/docker/cagent/pkg/tui/service"
+	"github.com/docker/cagent/pkg/tui/styles"
+	"github.com/docker/cagent/pkg/tui/types"
+)
+
+type Component struct {
+	message *types.Message
+	spinner spinner.Spinner
+	width   int
+	height  int
+}
+
+func New(
+	msg *types.Message,
+	_ *service.SessionState,
+) layout.Model {
+	return &Component{
+		message: msg,
+		spinner: spinner.New(spinner.ModeSpinnerOnly),
+		width:   80,
+		height:  1,
+	}
+}
+
+func (c *Component) SetSize(width, height int) tea.Cmd {
+	c.width = width
+	c.height = height
+	return nil
+}
+
+func (c *Component) Init() tea.Cmd {
+	if c.message.ToolStatus == types.ToolStatusPending || c.message.ToolStatus == types.ToolStatusRunning {
+		return c.spinner.Init()
+	}
+	return nil
+}
+
+func (c *Component) Update(msg tea.Msg) (layout.Model, tea.Cmd) {
+	if c.message.ToolStatus == types.ToolStatusPending || c.message.ToolStatus == types.ToolStatusRunning {
+		var cmd tea.Cmd
+		var model layout.Model
+		model, cmd = c.spinner.Update(msg)
+		c.spinner = model.(spinner.Spinner)
+		return c, cmd
+	}
+
+	return c, nil
+}
+
+func (c *Component) View() string {
+	msg := c.message
+	var args builtin.RunShellArgs
+	if err := json.Unmarshal([]byte(msg.ToolCall.Function.Arguments), &args); err != nil {
+		return toolcommon.RenderTool(toolcommon.Icon(msg.ToolStatus), msg.ToolDefinition.DisplayName(), c.spinner.View(), "", c.width)
+	}
+
+	return toolcommon.RenderTool(toolcommon.Icon(msg.ToolStatus), msg.ToolDefinition.DisplayName(), styles.MutedStyle.Render(args.Cmd), "", c.width)
+}