chore: rework layout primitives

7 files changed, 274 insertions, 565 deletions

diff --git a/packages/tui/internal/components/chat/editor.go b/packages/tui/internal/components/chat/editor.go
index 0ac3978a8..57dcbf154 100644
--- a/packages/tui/internal/components/chat/editor.go
+++ b/packages/tui/internal/components/chat/editor.go
@@ -22,7 +22,7 @@ import (
 type EditorComponent interface {
 	tea.Model
 	tea.ViewModel
-	layout.Sizeable
+	SetSize(width, height int) tea.Cmd
 	Content() string
 	Lines() int
 	Value() string
@@ -158,7 +158,15 @@ func (m *editorComponent) Content() string {
 
 func (m *editorComponent) View() string {
 	if m.Lines() > 1 {
-		return ""
+		t := theme.CurrentTheme()
+		return lipgloss.Place(
+			m.width,
+			m.height,
+			lipgloss.Center,
+			lipgloss.Center,
+			"",
+			styles.WhitespaceStyle(t.Background()),
+		)
 	}
 	return m.Content()
 }
diff --git a/packages/tui/internal/components/chat/messages.go b/packages/tui/internal/components/chat/messages.go
index 0a1aaa8f8..ab9ef65ec 100644
--- a/packages/tui/internal/components/chat/messages.go
+++ b/packages/tui/internal/components/chat/messages.go
@@ -21,6 +21,7 @@ import (
 type MessagesComponent interface {
 	tea.Model
 	tea.ViewModel
+	SetSize(width, height int) tea.Cmd
 	PageUp() (tea.Model, tea.Cmd)
 	PageDown() (tea.Model, tea.Cmd)
 	HalfPageUp() (tea.Model, tea.Cmd)
@@ -311,6 +312,7 @@ func (m *messagesComponent) View() string {
 	if len(m.app.Messages) == 0 {
 		return m.home()
 	}
+	t := theme.CurrentTheme()
 	if m.rendering {
 		return lipgloss.Place(
 			m.width,
@@ -318,19 +320,18 @@ func (m *messagesComponent) View() string {
 			lipgloss.Center,
 			lipgloss.Center,
 			"Loading session...",
+			styles.WhitespaceStyle(t.Background()),
 		)
 	}
-	t := theme.CurrentTheme()
-	return lipgloss.JoinVertical(
-		lipgloss.Left,
-		lipgloss.PlaceHorizontal(
-			m.width,
-			lipgloss.Center,
-			m.header(),
-			styles.WhitespaceStyle(t.Background()),
-		),
-		m.viewport.View(),
+	header := lipgloss.PlaceHorizontal(
+		m.width,
+		lipgloss.Center,
+		m.header(),
+		styles.WhitespaceStyle(t.Background()),
 	)
+	return styles.NewStyle().
+		Background(t.Background()).
+		Render(header + "\n" + m.viewport.View())
 }
 
 func (m *messagesComponent) home() string {
diff --git a/packages/tui/internal/components/commands/commands.go b/packages/tui/internal/components/commands/commands.go
index e20755575..68f6503e0 100644
--- a/packages/tui/internal/components/commands/commands.go
+++ b/packages/tui/internal/components/commands/commands.go
@@ -9,7 +9,6 @@ import (
 	"github.com/charmbracelet/lipgloss/v2/compat"
 	"github.com/sst/opencode/internal/app"
 	"github.com/sst/opencode/internal/commands"
-	"github.com/sst/opencode/internal/layout"
 	"github.com/sst/opencode/internal/styles"
 	"github.com/sst/opencode/internal/theme"
 )
@@ -17,7 +16,7 @@ import (
 type CommandsComponent interface {
 	tea.Model
 	tea.ViewModel
-	layout.Sizeable
+	SetSize(width, height int) tea.Cmd
 	SetBackgroundColor(color compat.AdaptiveColor)
 }
 
diff --git a/packages/tui/internal/layout/container.go b/packages/tui/internal/layout/container.go
deleted file mode 100644
index 250034ebc..000000000
--- a/packages/tui/internal/layout/container.go
+++ /dev/null
@@ -1,292 +0,0 @@
-package layout
-
-import (
-	tea "github.com/charmbracelet/bubbletea/v2"
-	"github.com/charmbracelet/lipgloss/v2"
-	"github.com/sst/opencode/internal/styles"
-	"github.com/sst/opencode/internal/theme"
-)
-
-type Container interface {
-	tea.Model
-	tea.ViewModel
-	Sizeable
-	Focusable
-	Alignable
-}
-
-type container struct {
-	width  int
-	height int
-	x      int
-	y      int
-
-	content tea.ViewModel
-
-	paddingTop    int
-	paddingRight  int
-	paddingBottom int
-	paddingLeft   int
-
-	borderTop    bool
-	borderRight  bool
-	borderBottom bool
-	borderLeft   bool
-	borderStyle  lipgloss.Border
-
-	maxWidth int
-	align    lipgloss.Position
-
-	focused bool
-}
-
-func (c *container) Init() tea.Cmd {
-	if model, ok := c.content.(tea.Model); ok {
-		return model.Init()
-	}
-	return nil
-}
-
-func (c *container) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
-	if model, ok := c.content.(tea.Model); ok {
-		u, cmd := model.Update(msg)
-		c.content = u.(tea.ViewModel)
-		return c, cmd
-	}
-	return c, nil
-}
-
-func (c *container) View() string {
-	t := theme.CurrentTheme()
-	style := styles.NewStyle().Background(t.Background())
-	width := c.width
-	height := c.height
-
-	// Apply max width constraint if set
-	if c.maxWidth > 0 && width > c.maxWidth {
-		width = c.maxWidth
-	}
-
-	// Apply border if any side is enabled
-	if c.borderTop || c.borderRight || c.borderBottom || c.borderLeft {
-		// Adjust width and height for borders
-		if c.borderTop {
-			height--
-		}
-		if c.borderBottom {
-			height--
-		}
-		if c.borderLeft {
-			width--
-		}
-		if c.borderRight {
-			width--
-		}
-		style = style.Border(c.borderStyle, c.borderTop, c.borderRight, c.borderBottom, c.borderLeft)
-
-		// Use primary color for border if focused
-		if c.focused {
-			style = style.BorderBackground(t.Background()).BorderForeground(t.Primary())
-		} else {
-			style = style.BorderBackground(t.Background()).BorderForeground(t.Border())
-		}
-	}
-	style = style.
-		Width(width).
-		Height(height).
-		PaddingTop(c.paddingTop).
-		PaddingRight(c.paddingRight).
-		PaddingBottom(c.paddingBottom).
-		PaddingLeft(c.paddingLeft)
-
-	return style.Render(c.content.View())
-}
-
-func (c *container) SetSize(width, height int) tea.Cmd {
-	c.width = width
-	c.height = height
-
-	// Apply max width constraint if set
-	effectiveWidth := width
-	if c.maxWidth > 0 && width > c.maxWidth {
-		effectiveWidth = c.maxWidth
-	}
-
-	// If the content implements Sizeable, adjust its size to account for padding and borders
-	if sizeable, ok := c.content.(Sizeable); ok {
-		// Calculate horizontal space taken by padding and borders
-		horizontalSpace := c.paddingLeft + c.paddingRight
-		if c.borderLeft {
-			horizontalSpace++
-		}
-		if c.borderRight {
-			horizontalSpace++
-		}
-
-		// Calculate vertical space taken by padding and borders
-		verticalSpace := c.paddingTop + c.paddingBottom
-		if c.borderTop {
-			verticalSpace++
-		}
-		if c.borderBottom {
-			verticalSpace++
-		}
-
-		// Set content size with adjusted dimensions
-		contentWidth := max(0, effectiveWidth-horizontalSpace)
-		contentHeight := max(0, height-verticalSpace)
-		return sizeable.SetSize(contentWidth, contentHeight)
-	}
-	return nil
-}
-
-func (c *container) GetSize() (int, int) {
-	return min(c.width, c.maxWidth), c.height
-}
-
-func (c *container) MaxWidth() int {
-	return c.maxWidth
-}
-
-func (c *container) Alignment() lipgloss.Position {
-	return c.align
-}
-
-// Focus sets the container as focused
-func (c *container) Focus() tea.Cmd {
-	c.focused = true
-	if focusable, ok := c.content.(Focusable); ok {
-		return focusable.Focus()
-	}
-	return nil
-}
-
-// Blur removes focus from the container
-func (c *container) Blur() tea.Cmd {
-	c.focused = false
-	if blurable, ok := c.content.(Focusable); ok {
-		return blurable.Blur()
-	}
-	return nil
-}
-
-func (c *container) IsFocused() bool {
-	if blurable, ok := c.content.(Focusable); ok {
-		return blurable.IsFocused()
-	}
-	return c.focused
-}
-
-// GetPosition returns the x, y coordinates of the container
-func (c *container) GetPosition() (x, y int) {
-	return c.x, c.y
-}
-
-func (c *container) SetPosition(x, y int) {
-	c.x = x
-	c.y = y
-}
-
-type ContainerOption func(*container)
-
-func NewContainer(content tea.ViewModel, options ...ContainerOption) Container {
-	c := &container{
-		content:     content,
-		borderStyle: lipgloss.NormalBorder(),
-	}
-	for _, option := range options {
-		option(c)
-	}
-	return c
-}
-
-// Padding options
-func WithPadding(top, right, bottom, left int) ContainerOption {
-	return func(c *container) {
-		c.paddingTop = top
-		c.paddingRight = right
-		c.paddingBottom = bottom
-		c.paddingLeft = left
-	}
-}
-
-func WithPaddingAll(padding int) ContainerOption {
-	return WithPadding(padding, padding, padding, padding)
-}
-
-func WithPaddingHorizontal(padding int) ContainerOption {
-	return func(c *container) {
-		c.paddingLeft = padding
-		c.paddingRight = padding
-	}
-}
-
-func WithPaddingVertical(padding int) ContainerOption {
-	return func(c *container) {
-		c.paddingTop = padding
-		c.paddingBottom = padding
-	}
-}
-
-func WithBorder(top, right, bottom, left bool) ContainerOption {
-	return func(c *container) {
-		c.borderTop = top
-		c.borderRight = right
-		c.borderBottom = bottom
-		c.borderLeft = left
-	}
-}
-
-func WithBorderAll() ContainerOption {
-	return WithBorder(true, true, true, true)
-}
-
-func WithBorderHorizontal() ContainerOption {
-	return WithBorder(true, false, true, false)
-}
-
-func WithBorderVertical() ContainerOption {
-	return WithBorder(false, true, false, true)
-}
-
-func WithBorderStyle(style lipgloss.Border) ContainerOption {
-	return func(c *container) {
-		c.borderStyle = style
-	}
-}
-
-func WithRoundedBorder() ContainerOption {
-	return WithBorderStyle(lipgloss.RoundedBorder())
-}
-
-func WithThickBorder() ContainerOption {
-	return WithBorderStyle(lipgloss.ThickBorder())
-}
-
-func WithDoubleBorder() ContainerOption {
-	return WithBorderStyle(lipgloss.DoubleBorder())
-}
-
-func WithMaxWidth(maxWidth int) ContainerOption {
-	return func(c *container) {
-		c.maxWidth = maxWidth
-	}
-}
-
-func WithAlign(align lipgloss.Position) ContainerOption {
-	return func(c *container) {
-		c.align = align
-	}
-}
-
-func WithAlignLeft() ContainerOption {
-	return WithAlign(lipgloss.Left)
-}
-
-func WithAlignCenter() ContainerOption {
-	return WithAlign(lipgloss.Center)
-}
-
-func WithAlignRight() ContainerOption {
-	return WithAlign(lipgloss.Right)
-}
diff --git a/packages/tui/internal/layout/flex.go b/packages/tui/internal/layout/flex.go
index 320a95203..f164a03dc 100644
--- a/packages/tui/internal/layout/flex.go
+++ b/packages/tui/internal/layout/flex.go
@@ -1,255 +1,260 @@
 package layout
 
 import (
-	tea "github.com/charmbracelet/bubbletea/v2"
+	"strings"
+
 	"github.com/charmbracelet/lipgloss/v2"
 	"github.com/sst/opencode/internal/styles"
-	"github.com/sst/opencode/internal/theme"
 )
 
-type FlexDirection int
+type Direction int
 
 const (
-	FlexDirectionHorizontal FlexDirection = iota
-	FlexDirectionVertical
+	Row Direction = iota
+	Column
 )
 
-type FlexChildSize struct {
-	Fixed bool
-	Size  int
-}
+type Justify int
 
-var FlexChildSizeGrow = FlexChildSize{Fixed: false}
+const (
+	JustifyStart Justify = iota
+	JustifyEnd
+	JustifyCenter
+	JustifySpaceBetween
+	JustifySpaceAround
+)
 
-func FlexChildSizeFixed(size int) FlexChildSize {
-	return FlexChildSize{Fixed: true, Size: size}
-}
+type Align int
 
-type FlexLayout interface {
-	tea.ViewModel
-	Sizeable
-	SetChildren(panes []tea.ViewModel) tea.Cmd
-	SetSizes(sizes []FlexChildSize) tea.Cmd
-	SetDirection(direction FlexDirection) tea.Cmd
-}
+const (
+	AlignStart Align = iota
+	AlignEnd
+	AlignCenter
+	AlignStretch // Only applicable in the cross-axis
+)
 
-type flexLayout struct {
-	width     int
-	height    int
-	direction FlexDirection
-	children  []tea.ViewModel
-	sizes     []FlexChildSize
+type FlexOptions struct {
+	Direction Direction
+	Justify   Justify
+	Align     Align
+	Width     int
+	Height    int
 }
 
-type FlexLayoutOption func(*flexLayout)
+type FlexItem struct {
+	View      string
+	FixedSize int  // Fixed size in the main axis (width for Row, height for Column)
+	Grow      bool // If true, the item will grow to fill available space
+}
 
-func (f *flexLayout) View() string {
-	if len(f.children) == 0 {
+// Render lays out a series of view strings based on flexbox-like rules.
+func Render(opts FlexOptions, items ...FlexItem) string {
+	if len(items) == 0 {
 		return ""
 	}
 
-	t := theme.CurrentTheme()
-	views := make([]string, 0, len(f.children))
-	for i, child := range f.children {
-		if child == nil {
-			continue
-		}
+	// Calculate dimensions for each item
+	mainAxisSize := opts.Width
+	crossAxisSize := opts.Height
+	if opts.Direction == Column {
+		mainAxisSize = opts.Height
+		crossAxisSize = opts.Width
+	}
 
-		alignment := lipgloss.Center
-		if alignable, ok := child.(Alignable); ok {
-			alignment = alignable.Alignment()
-		}
-		var childWidth, childHeight int
-		if f.direction == FlexDirectionHorizontal {
-			childWidth, childHeight = f.calculateChildSize(i)
-			view := lipgloss.PlaceHorizontal(
-				childWidth,
-				alignment,
-				child.View(),
-				// TODO: make configurable WithBackgroundStyle
-				lipgloss.WithWhitespaceStyle(styles.NewStyle().Background(t.Background()).Lipgloss()),
-			)
-			views = append(views, view)
-		} else {
-			childWidth, childHeight = f.calculateChildSize(i)
-			view := lipgloss.Place(
-				f.width,
-				childHeight,
-				lipgloss.Center,
-				alignment,
-				child.View(),
-				// TODO: make configurable WithBackgroundStyle
-				lipgloss.WithWhitespaceStyle(styles.NewStyle().Background(t.Background()).Lipgloss()),
-			)
-			views = append(views, view)
+	// Calculate total fixed size and count grow items
+	totalFixedSize := 0
+	growCount := 0
+	for _, item := range items {
+		if item.FixedSize > 0 {
+			totalFixedSize += item.FixedSize
+		} else if item.Grow {
+			growCount++
 		}
 	}
-	if f.direction == FlexDirectionHorizontal {
-		return lipgloss.JoinHorizontal(lipgloss.Center, views...)
+
+	// Calculate available space for grow items
+	availableSpace := mainAxisSize - totalFixedSize
+	if availableSpace < 0 {
+		availableSpace = 0
 	}
-	return lipgloss.JoinVertical(lipgloss.Center, views...)
-}
 
-func (f *flexLayout) calculateChildSize(index int) (width, height int) {
-	if index >= len(f.children) {
-		return 0, 0
+	// Calculate size for each grow item
+	growItemSize := 0
+	if growCount > 0 && availableSpace > 0 {
+		growItemSize = availableSpace / growCount
 	}
 
-	totalFixed := 0
-	flexCount := 0
+	// Prepare sized views
+	sizedViews := make([]string, len(items))
+	actualSizes := make([]int, len(items))
 
-	for i, child := range f.children {
-		if child == nil {
-			continue
-		}
-		if i < len(f.sizes) && f.sizes[i].Fixed {
-			if f.direction == FlexDirectionHorizontal {
-				totalFixed += f.sizes[i].Size
+	for i, item := range items {
+		view := item.View
+
+		// Determine the size for this item
+		itemSize := 0
+		if item.FixedSize > 0 {
+			itemSize = item.FixedSize
+		} else if item.Grow && growItemSize > 0 {
+			itemSize = growItemSize
+		} else {
+			// No fixed size and not growing - use natural size
+			if opts.Direction == Row {
+				itemSize = lipgloss.Width(view)
 			} else {
-				totalFixed += f.sizes[i].Size
+				itemSize = lipgloss.Height(view)
 			}
-		} else {
-			flexCount++
 		}
-	}
 
-	if f.direction == FlexDirectionHorizontal {
-		height = f.height
-		if index < len(f.sizes) && f.sizes[index].Fixed {
-			width = f.sizes[index].Size
-		} else if flexCount > 0 {
-			remainingSpace := f.width - totalFixed
-			width = remainingSpace / flexCount
-		}
-	} else {
-		width = f.width
-		if index < len(f.sizes) && f.sizes[index].Fixed {
-			height = f.sizes[index].Size
-		} else if flexCount > 0 {
-			remainingSpace := f.height - totalFixed
-			height = remainingSpace / flexCount
-		}
-	}
-
-	return width, height
-}
-
-func (f *flexLayout) SetSize(width, height int) tea.Cmd {
-	f.width = width
-	f.height = height
-
-	var cmds []tea.Cmd
-	currentX, currentY := 0, 0
-
-	for i, child := range f.children {
-		if child != nil {
-			paneWidth, paneHeight := f.calculateChildSize(i)
-			alignment := lipgloss.Center
-			if alignable, ok := child.(Alignable); ok {
-				alignment = alignable.Alignment()
+		// Apply size constraints
+		if opts.Direction == Row {
+			// For row direction, constrain width and handle height alignment
+			if itemSize > 0 {
+				view = styles.NewStyle().
+					Width(itemSize).
+					Height(crossAxisSize).
+					Render(view)
 			}
 
-			// Calculate actual position based on alignment
-			actualX, actualY := currentX, currentY
-
-			if f.direction == FlexDirectionHorizontal {
-				// In horizontal layout, vertical alignment affects Y position
-				// (lipgloss.Center is used for vertical alignment in JoinHorizontal)
-				actualY = (f.height - paneHeight) / 2
-			} else {
-				// In vertical layout, horizontal alignment affects X position
-				contentWidth := paneWidth
-				if alignable, ok := child.(Alignable); ok {
-					if alignable.MaxWidth() > 0 && contentWidth > alignable.MaxWidth() {
-						contentWidth = alignable.MaxWidth()
-					}
-				}
-
-				switch alignment {
-				case lipgloss.Center:
-					actualX = (f.width - contentWidth) / 2
-				case lipgloss.Right:
-					actualX = f.width - contentWidth
-				case lipgloss.Left:
-					actualX = 0
-				}
+			// Apply cross-axis alignment
+			switch opts.Align {
+			case AlignCenter:
+				view = lipgloss.PlaceVertical(crossAxisSize, lipgloss.Center, view)
+			case AlignEnd:
+				view = lipgloss.PlaceVertical(crossAxisSize, lipgloss.Bottom, view)
+			case AlignStart:
+				view = lipgloss.PlaceVertical(crossAxisSize, lipgloss.Top, view)
+			case AlignStretch:
+				// Already stretched by Height setting above
 			}
-
-			// Set position if the pane is Alignable
-			if c, ok := child.(Alignable); ok {
-				c.SetPosition(actualX, actualY)
+		} else {
+			// For column direction, constrain height and handle width alignment
+			if itemSize > 0 {
+				view = styles.NewStyle().
+					Height(itemSize).
+					Width(crossAxisSize).
+					Render(view)
 			}
 
-			if sizeable, ok := child.(Sizeable); ok {
-				cmd := sizeable.SetSize(paneWidth, paneHeight)
-				cmds = append(cmds, cmd)
+			// Apply cross-axis alignment
+			switch opts.Align {
+			case AlignCenter:
+				view = lipgloss.PlaceHorizontal(crossAxisSize, lipgloss.Center, view)
+			case AlignEnd:
+				view = lipgloss.PlaceHorizontal(crossAxisSize, lipgloss.Right, view)
+			case AlignStart:
+				view = lipgloss.PlaceHorizontal(crossAxisSize, lipgloss.Left, view)
+			case AlignStretch:
+				// Already stretched by Width setting above
 			}
+		}
 
-			// Update position for next pane
-			if f.direction == FlexDirectionHorizontal {
-				currentX += paneWidth
-			} else {
-				currentY += paneHeight
-			}
+		sizedViews[i] = view
+		if opts.Direction == Row {
+			actualSizes[i] = lipgloss.Width(view)
+		} else {
+			actualSizes[i] = lipgloss.Height(view)
 		}
 	}
-	return tea.Batch(cmds...)
-}
 
-func (f *flexLayout) GetSize() (int, int) {
-	return f.width, f.height
-}
+	// Calculate total actual size
+	totalActualSize := 0
+	for _, size := range actualSizes {
+		totalActualSize += size
+	}
 
-func (f *flexLayout) SetChildren(children []tea.ViewModel) tea.Cmd {
-	f.children = children
-	if f.width > 0 && f.height > 0 {
-		return f.SetSize(f.width, f.height)
+	// Apply justification
+	remainingSpace := mainAxisSize - totalActualSize
+	if remainingSpace < 0 {
+		remainingSpace = 0
 	}
-	return nil
-}
 
-func (f *flexLayout) SetSizes(sizes []FlexChildSize) tea.Cmd {
-	f.sizes = sizes
-	if f.width > 0 && f.height > 0 {
-		return f.SetSize(f.width, f.height)
+	// Calculate spacing based on justification
+	var spaceBefore, spaceBetween, spaceAfter int
+	switch opts.Justify {
+	case JustifyStart:
+		spaceAfter = remainingSpace
+	case JustifyEnd:
+		spaceBefore = remainingSpace
+	case JustifyCenter:
+		spaceBefore = remainingSpace / 2
+		spaceAfter = remainingSpace - spaceBefore
+	case JustifySpaceBetween:
+		if len(items) > 1 {
+			spaceBetween = remainingSpace / (len(items) - 1)
+		} else {
+			spaceAfter = remainingSpace
+		}
+	case JustifySpaceAround:
+		if len(items) > 0 {
+			spaceAround := remainingSpace / (len(items) * 2)
+			spaceBefore = spaceAround
+			spaceAfter = spaceAround
+			spaceBetween = spaceAround * 2
+		}
 	}
-	return nil
-}
 
-func (f *flexLayout) SetDirection(direction FlexDirection) tea.Cmd {
-	f.direction = direction
-	if f.width > 0 && f.height > 0 {
-		return f.SetSize(f.width, f.height)
+	// Build the final layout
+	var parts []string
+
+	// Add space before if needed
+	if spaceBefore > 0 {
+		if opts.Direction == Row {
+			parts = append(parts, strings.Repeat(" ", spaceBefore))
+		} else {
+			parts = append(parts, strings.Repeat("\n", spaceBefore))
+		}
 	}
-	return nil
-}
 
-func NewFlexLayout(children []tea.ViewModel, options ...FlexLayoutOption) FlexLayout {
-	layout := &flexLayout{
-		children:  children,
-		direction: FlexDirectionHorizontal,
-		sizes:     []FlexChildSize{},
+	// Add items with spacing
+	for i, view := range sizedViews {
+		parts = append(parts, view)
+
+		// Add space between items (not after the last one)
+		if i < len(sizedViews)-1 && spaceBetween > 0 {
+			if opts.Direction == Row {
+				parts = append(parts, strings.Repeat(" ", spaceBetween))
+			} else {
+				parts = append(parts, strings.Repeat("\n", spaceBetween))
+			}
+		}
 	}
-	for _, option := range options {
-		option(layout)
+
+	// Add space after if needed
+	if spaceAfter > 0 {
+		if opts.Direction == Row {
+			parts = append(parts, strings.Repeat(" ", spaceAfter))
+		} else {
+			parts = append(parts, strings.Repeat("\n", spaceAfter))
+		}
 	}
-	return layout
-}
 
-func WithDirection(direction FlexDirection) FlexLayoutOption {
-	return func(f *flexLayout) {
-		f.direction = direction
+	// Join the parts
+	if opts.Direction == Row {
+		return lipgloss.JoinHorizontal(lipgloss.Top, parts...)
+	} else {
+		return lipgloss.JoinVertical(lipgloss.Left, parts...)
 	}
 }
 
-func WithChildren(children ...tea.ViewModel) FlexLayoutOption {
-	return func(f *flexLayout) {
-		f.children = children
-	}
+// Helper function to create a simple vertical layout
+func Vertical(width, height int, items ...FlexItem) string {
+	return Render(FlexOptions{
+		Direction: Column,
+		Width:     width,
+		Height:    height,
+		Justify:   JustifyStart,
+		Align:     AlignStretch,
+	}, items...)
 }
 
-func WithSizes(sizes ...FlexChildSize) FlexLayoutOption {
-	return func(f *flexLayout) {
-		f.sizes = sizes
-	}
+// Helper function to create a simple horizontal layout
+func Horizontal(width, height int, items ...FlexItem) string {
+	return Render(FlexOptions{
+		Direction: Row,
+		Width:     width,
+		Height:    height,
+		Justify:   JustifyStart,
+		Align:     AlignStretch,
+	}, items...)
 }
diff --git a/packages/tui/internal/layout/layout.go b/packages/tui/internal/layout/layout.go
index 208faaa2f..dce27ac68 100644
--- a/packages/tui/internal/layout/layout.go
+++ b/packages/tui/internal/layout/layout.go
@@ -1,11 +1,7 @@
 package layout
 
 import (
-	"reflect"
-
-	"github.com/charmbracelet/bubbles/v2/key"
 	tea "github.com/charmbracelet/bubbletea/v2"
-	"github.com/charmbracelet/lipgloss/v2"
 )
 
 var Current *LayoutInfo
@@ -34,33 +30,3 @@ type Modal interface {
 	Render(background string) string
 	Close() tea.Cmd
 }
-
-type Focusable interface {
-	Focus() tea.Cmd
-	Blur() tea.Cmd
-	IsFocused() bool
-}
-
-type Sizeable interface {
-	SetSize(width, height int) tea.Cmd
-	GetSize() (int, int)
-}
-
-type Alignable interface {
-	MaxWidth() int
-	Alignment() lipgloss.Position
-	SetPosition(x, y int)
-	GetPosition() (x, y int)
-}
-
-func KeyMapToSlice(t any) (bindings []key.Binding) {
-	typ := reflect.TypeOf(t)
-	if typ.Kind() != reflect.Struct {
-		return nil
-	}
-	for i := range typ.NumField() {
-		v := reflect.ValueOf(t).Field(i)
-		bindings = append(bindings, v.Interface().(key.Binding))
-	}
-	return
-}
diff --git a/packages/tui/internal/tui/tui.go b/packages/tui/internal/tui/tui.go
index 28ce9f282..67538e804 100644
--- a/packages/tui/internal/tui/tui.go
+++ b/packages/tui/internal/tui/tui.go
@@ -47,8 +47,6 @@ type appModel struct {
 	status               status.StatusComponent
 	editor               chat.EditorComponent
 	messages             chat.MessagesComponent
-	editorContainer      layout.Container
-	layout               layout.FlexLayout
 	completions          dialog.CompletionDialog
 	completionManager    *completions.CompletionManager
 	showCompletionDialog bool
@@ -360,7 +358,10 @@ func (a appModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 				Width: min(a.width, 80),
 			},
 		}
-		a.layout.SetSize(a.width, a.height)
+		// Update child component sizes
+		messagesHeight := a.height - 6 // Leave room for editor and status bar
+		a.messages.SetSize(a.width, messagesHeight)
+		a.editor.SetSize(min(a.width, 80), 5)
 	case app.SessionSelectedMsg:
 		messages, err := a.app.ListMessages(context.Background(), msg.ID)
 		if err != nil {
@@ -424,33 +425,69 @@ func (a appModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 }
 
 func (a appModel) View() string {
-	layoutView := a.layout.View()
-	editorWidth, _ := a.editorContainer.GetSize()
-	editorX, editorY := a.editorContainer.GetPosition()
+	messagesView := a.messages.View()
+	editorView := a.editor.View()
+
+	editorHeight := lipgloss.Height(editorView)
+	if editorHeight < 5 {
+		editorHeight = 5
+	}
+
+	t := theme.CurrentTheme()
+	centeredEditorView := lipgloss.PlaceHorizontal(
+		a.width,
+		lipgloss.Center,
+		editorView,
+		styles.WhitespaceStyle(t.Background()),
+	)
+
+	mainLayout := layout.Render(
+		layout.FlexOptions{
+			Direction: layout.Column,
+			Width:     a.width,
+			Height:    a.height - 1, // Leave room for status bar
+		},
+		layout.FlexItem{
+			View: messagesView,
+			Grow: true,
+		},
+		layout.FlexItem{
+			View:      centeredEditorView,
+			FixedSize: editorHeight,
+		},
+	)
 
 	if a.editor.Lines() > 1 {
-		editorY = editorY - a.editor.Lines() + 1
-		layoutView = layout.PlaceOverlay(
+		editorWidth := min(a.width, 80)
+		editorX := (a.width - editorWidth) / 2
+		editorY := a.height - editorHeight - 1 // Position from bottom, accounting for status bar
+
+		mainLayout = layout.PlaceOverlay(
 			editorX,
 			editorY,
 			a.editor.Content(),
-			layoutView,
+			mainLayout,
 		)
 	}
 
 	if a.showCompletionDialog {
+		editorWidth := min(a.width, 80)
+		editorX := (a.width - editorWidth) / 2
 		a.completions.SetWidth(editorWidth)
 		overlay := a.completions.View()
-		layoutView = layout.PlaceOverlay(
+		overlayHeight := lipgloss.Height(overlay)
+		editorY := a.height - editorHeight - 1
+
+		mainLayout = layout.PlaceOverlay(
 			editorX,
-			editorY-lipgloss.Height(overlay)+2,
+			editorY-overlayHeight,
 			overlay,
-			layoutView,
+			mainLayout,
 		)
 	}
 
 	components := []string{
-		layoutView,
+		mainLayout,
 		a.status.View(),
 	}
 	appView := strings.Join(components, "\n")
@@ -464,6 +501,7 @@ func (a appModel) View() string {
 	if theme.CurrentThemeUsesAnsiColors() {
 		appView = util.ConvertRGBToAnsi16Colors(appView)
 	}
+
 	return appView
 }
 
@@ -653,13 +691,6 @@ func NewModel(app *app.App) tea.Model {
 	editor := chat.NewEditorComponent(app)
 	completions := dialog.NewCompletionDialogComponent(initialProvider)
 
-	editorContainer := layout.NewContainer(
-		editor,
-		layout.WithMaxWidth(layout.Current.Container.Width),
-		layout.WithAlignCenter(),
-	)
-	messagesContainer := layout.NewContainer(messages)
-
 	var leaderBinding *key.Binding
 	if app.Config.Keybinds.Leader != "" {
 		binding := key.NewBinding(key.WithKeys(app.Config.Keybinds.Leader))
@@ -676,17 +707,8 @@ func NewModel(app *app.App) tea.Model {
 		leaderBinding:        leaderBinding,
 		isLeaderSequence:     false,
 		showCompletionDialog: false,
-		editorContainer:      editorContainer,
 		toastManager:         toast.NewToastManager(),
 		interruptKeyState:    InterruptKeyIdle,
-		layout: layout.NewFlexLayout(
-			[]tea.ViewModel{messagesContainer, editorContainer},
-			layout.WithDirection(layout.FlexDirectionVertical),
-			layout.WithSizes(
-				layout.FlexChildSizeGrow,
-				layout.FlexChildSizeFixed(5),
-			),
-		),
 	}
 
 	return model