Merge pull request #3486 from docker/feat/tui-per-agent-context

feat(tui): per-sub-agent context accounting in the TUI

Author: dgageot

docs/features/tui/index.md

@@ -100,7 +100,7 @@ Slash commands (both built-in and named) execute immediately when entered. Regul
 
 ### Agents Panel
 
-The sidebar's **Agents** section lists every agent in the team. The current agent is shown as a focus **card** (rendered in place at its position in the list) with its name, a wrapped description, its full `provider/model`, and a thinking line. Every other agent is shown as a compact **two-line row** — line 1 is the shortcut/spinner, the agent name (in its accent color), and a right-aligned thinking **gauge**; line 2 is the indented full `provider/model` — so a large team stays scannable while still showing each model. Agents are separated by a blank line so the two-line rows stay visually distinct. The effort **gauge** is the only visual language for thinking; the focus card and the Agent Inspector spell out the exact level alongside it. Left-click any agent to switch to it.
+The sidebar's **Agents** section lists every agent in the team. The current agent is shown as a focus **card** (rendered in place at its position in the list) with its name, a wrapped description, its full `provider/model`, and a thinking line. Every other agent is shown as a compact **two-line row** — line 1 is the shortcut/spinner, the agent name (in its accent color), and a right-aligned thinking **gauge**; line 2 is the indented full `provider/model` — so a large team stays scannable while still showing each model. Once an agent has run (in the main session, as a delegated sub-agent, or as a background agent task), line 2 also carries its latest **context usage** as a right-aligned percentage of its context window, so per-agent context accounting is visible at a glance across the whole team. Agents are separated by a blank line so the two-line rows stay visually distinct. The effort **gauge** is the only visual language for thinking; the focus card and the Agent Inspector spell out the exact level alongside it. Left-click any agent to switch to it.
 
 #### Agent inspector
 
@@ -115,6 +115,7 @@ The title is rendered in the agent's accent color. Sections appear in this order
 - **Description** — the agent's wrapped description.
 - **Live state** — a `● current agent` line when the inspected agent is the one currently running.
 - **Model / Fallback / Thinking** — the `provider/model`, any fallback models, and the gauge + value thinking line (omitted for models with no selectable thinking, e.g. harness-backed agents).
+- **Context** — the agent's latest known context usage, e.g. `Context: 12.8K of 128.0K tokens (10%)` (a bare token count when the context limit is unknown; omitted until the agent has run). Sub-agent and background-agent runs are accounted for.
 - **Sub-agents (N) / Handoffs (N) / Skills (N)** — compact, inline, comma-separated lists wrapped to the dialog width.
 - **Limits** — the configured per-agent limits that are set, e.g. `Limits: max-iter 50 · history 40 · max-tool-calls 5`.
 - **Options** — the enabled option flags, e.g. `Options: add-date · add-environment-info · redact-secrets`.

docs/tools/background-agents/index.md

@@ -76,6 +76,8 @@ agents:
 >
 > Use `background_agents` when your orchestrator needs to fan out work to multiple specialists in parallel — for example, researching several topics simultaneously or running independent code analyses side by side.
 
+In the TUI, each background task's token usage is accounted for live: the sidebar's Agents panel shows the sub-agent's context usage percentage on its roster row, the Agent Inspector shows its exact token counts, and the task's cost joins the session total.
+
 ## Using Harness Sub-Agents
 
 Background agents work equally well with [harness-backed sub-agents](../../features/harnesses/index.md) — sub-agents driven by external coding CLIs such as Claude Code or Codex. This lets you dispatch multiple independent coding tasks in parallel:

e2e/tui/background_agent_test.go

@@ -0,0 +1,214 @@
+package tui_test
+
+import (
+	"context"
+	"io"
+	"os"
+	"sync"
+	"testing"
+
+	tea "charm.land/bubbletea/v2"
+	"github.com/stretchr/testify/require"
+
+	"github.com/docker/docker-agent/pkg/agent"
+	"github.com/docker/docker-agent/pkg/app"
+	"github.com/docker/docker-agent/pkg/chat"
+	"github.com/docker/docker-agent/pkg/config/latest"
+	"github.com/docker/docker-agent/pkg/model/provider/base"
+	"github.com/docker/docker-agent/pkg/modelsdev"
+	"github.com/docker/docker-agent/pkg/runtime"
+	"github.com/docker/docker-agent/pkg/session"
+	"github.com/docker/docker-agent/pkg/team"
+	"github.com/docker/docker-agent/pkg/tools"
+	agenttool "github.com/docker/docker-agent/pkg/tools/builtin/agent"
+	"github.com/docker/docker-agent/pkg/tui"
+	"github.com/docker/docker-agent/pkg/tui/tuitest"
+)
+
+// This file covers the background-agent journey end to end through the real
+// TUI: a root agent dispatches a run_background_agent task, the worker's
+// sub-session runs on a detached goroutine (its events are dropped from the
+// live stream), and only the runtime's out-of-band forwarding can bring its
+// token usage back to the sidebar. Unlike the VCR-based scenarios in this
+// package, the model responses are scripted in-process: driving a
+// deterministic multi-agent tool-call flow through cassettes would couple the
+// test to HTTP wire details it does not care about.
+
+// scriptedStream replays a fixed sequence of streaming chunks.
+type scriptedStream struct {
+	responses []chat.MessageStreamResponse
+	idx       int
+}
+
+func (s *scriptedStream) Recv() (chat.MessageStreamResponse, error) {
+	if s.idx >= len(s.responses) {
+		return chat.MessageStreamResponse{}, io.EOF
+	}
+	resp := s.responses[s.idx]
+	s.idx++
+	return resp, nil
+}
+
+func (s *scriptedStream) Close() {}
+
+// scriptedProvider returns one scripted stream per model call, repeating the
+// last script if called more often. contextSize is exposed through
+// provider_opts so the runtime resolves a context limit without the
+// models.dev catalogue, which is what makes percentages computable.
+type scriptedProvider struct {
+	id          string
+	contextSize int64
+	scripts     [][]chat.MessageStreamResponse
+
+	mu   sync.Mutex
+	call int
+}
+
+func (p *scriptedProvider) ID() modelsdev.ID { return modelsdev.ParseIDOrZero(p.id) }
+
+func (p *scriptedProvider) CreateChatCompletionStream(context.Context, []chat.Message, []tools.Tool) (chat.MessageStream, error) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	i := min(p.call, len(p.scripts)-1)
+	p.call++
+	return &scriptedStream{responses: p.scripts[i]}, nil
+}
+
+func (p *scriptedProvider) BaseConfig() base.Config {
+	return base.Config{ModelConfig: latest.ModelConfig{
+		ProviderOpts: map[string]any{"context_size": p.contextSize},
+	}}
+}
+
+// toolCallScript scripts one turn that calls a single tool and stops with the
+// given usage.
+func toolCallScript(callID, toolName, args string, in, out int64) []chat.MessageStreamResponse {
+	return []chat.MessageStreamResponse{
+		{Choices: []chat.MessageStreamChoice{{Delta: chat.MessageDelta{ToolCalls: []tools.ToolCall{{
+			ID: callID, Type: "function", Function: tools.FunctionCall{Name: toolName},
+		}}}}}},
+		{Choices: []chat.MessageStreamChoice{{Delta: chat.MessageDelta{ToolCalls: []tools.ToolCall{{
+			ID: callID, Type: "function", Function: tools.FunctionCall{Arguments: args},
+		}}}}}},
+		{
+			Choices: []chat.MessageStreamChoice{{FinishReason: chat.FinishReasonToolCalls}},
+			Usage:   &chat.Usage{InputTokens: in, OutputTokens: out},
+		},
+	}
+}
+
+// contentScript scripts one plain-content turn that stops with the given usage.
+func contentScript(content string, in, out int64) []chat.MessageStreamResponse {
+	return []chat.MessageStreamResponse{
+		{Choices: []chat.MessageStreamChoice{{Delta: chat.MessageDelta{Content: content}}}},
+		{
+			Choices: []chat.MessageStreamChoice{{FinishReason: chat.FinishReasonStop}},
+			Usage:   &chat.Usage{InputTokens: in, OutputTokens: out},
+		},
+	}
+}
+
+// stubModelStore keeps the runtime off the network: no models.dev lookups, no
+// pricing. Context limits come from provider_opts instead.
+type stubModelStore struct{}
+
+func (stubModelStore) GetModel(context.Context, modelsdev.ID) (*modelsdev.Model, error) {
+	return nil, nil
+}
+func (stubModelStore) GetDatabase(context.Context) (*modelsdev.Database, error) { return nil, nil }
+
+// newBackgroundAgentTUI builds the real TUI over a real LocalRuntime whose
+// team is assembled programmatically: a root orchestrator with the
+// background_agents toolset and a worker sub-agent, both backed by scripted
+// providers. Tools are pre-approved so the dispatch needs no confirmation
+// dialog.
+func newBackgroundAgentTUI(t *testing.T, width, height int) *tuitest.Driver {
+	t.Helper()
+
+	isolateState(t)
+
+	rootProv := &scriptedProvider{
+		id:          "test/fake-root",
+		contextSize: 1000,
+		scripts: [][]chat.MessageStreamResponse{
+			toolCallScript("call-1", agenttool.ToolNameRunBackgroundAgent,
+				`{"agent":"worker","task":"count the files"}`, 40, 10),
+			// Session context tracks the LAST call's usage (its input already
+			// spans the whole prompt), so the root settles at 100/1000 = 10%.
+			contentScript("Dispatched the worker.", 80, 20),
+		},
+	}
+	workerProv := &scriptedProvider{
+		id:          "test/fake-worker",
+		contextSize: 1000,
+		scripts: [][]chat.MessageStreamResponse{
+			contentScript("worker finished the count", 300, 150),
+		},
+	}
+
+	worker := agent.New("worker", "Count things when asked.",
+		agent.WithModel(workerProv),
+		agent.WithDescription("Background worker"),
+	)
+	root := agent.New("root", "Dispatch work to the worker in the background.",
+		agent.WithModel(rootProv),
+		agent.WithDescription("Orchestrator"),
+		agent.WithSubAgents(worker),
+		agent.WithToolSets(agenttool.New()),
+	)
+
+	rt, err := runtime.New(t.Context(), team.New(team.WithAgents(root, worker)),
+		runtime.WithCurrentAgent("root"),
+		runtime.WithSessionCompaction(false),
+		runtime.WithModelStore(stubModelStore{}),
+	)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = rt.Close() })
+
+	application := app.New(t.Context(), rt, session.New(session.WithToolsApproved(true)))
+
+	wd, _ := os.Getwd()
+	model := tui.New(t.Context(), nil /* no spawner: single tab */, application, wd, func() {})
+	return tuitest.New(t, model, width, height)
+}
+
+// TestBackgroundAgent_PerAgentContextInSidebar drives the full journey: the
+// worker's usage (450 of 1000 tokens) can only reach the TUI through the
+// runtime's out-of-band background-event forwarding, so the 45% appearing on
+// the worker's roster line proves the whole chain — tool call, detached
+// sub-session, OnBackgroundEvent, app event stream, sidebar accounting.
+// The root's own line independently settles at 10% (its final turn consumed
+// 100 of 1000 tokens).
+func TestBackgroundAgent_PerAgentContextInSidebar(t *testing.T) {
+	d := newBackgroundAgentTUI(t, 120, 40)
+
+	// Both agents are listed before anything runs, with no context percent.
+	d.WaitFor(tuitest.ContainsAll("root", "worker")).
+		Assert(tuitest.Absent("%"))
+
+	d.Type("Please dispatch the worker.").
+		Enter().
+		WaitFor(tuitest.Contains("Dispatched the worker.")).
+		WaitFor(tuitest.Contains("45%")).
+		WaitFor(tuitest.Contains("10%"))
+}
+
+// TestBackgroundAgent_InspectorShowsWorkerContext opens the worker's Agent
+// Inspector (right-click on its roster line) after the background task
+// completed and checks the exact per-agent context line.
+func TestBackgroundAgent_InspectorShowsWorkerContext(t *testing.T) {
+	d := newBackgroundAgentTUI(t, 120, 40)
+
+	d.Type("Please dispatch the worker.").
+		Enter().
+		WaitFor(tuitest.Contains("Dispatched the worker.")).
+		WaitFor(tuitest.Contains("45%"))
+
+	// "45%" renders only on the worker's roster line, so its coordinates are
+	// a reliable right-click target for that agent (either of the entry's two
+	// lines opens the inspector).
+	x, y := d.MustFindText("45%")
+	d.Send(tea.MouseClickMsg{X: x, Y: y, Button: tea.MouseRight}).
+		Send(tea.MouseReleaseMsg{X: x, Y: y, Button: tea.MouseRight}).
+		WaitFor(tuitest.Contains("Context: 450 of 1.0K tokens (45%)"))
+}

pkg/app/app.go

@@ -168,6 +168,16 @@ func (a *App) Start(ctx context.Context) {
 			case <-ctx.Done():
 			}
 		})
+
+		// Forward events surfaced from detached background work (token usage
+		// from background agent tasks) so the sidebar and agent inspector can
+		// account for background agents' context usage.
+		a.runtime.OnBackgroundEvent(func(event runtime.Event) {
+			select {
+			case a.events <- event:
+			case <-ctx.Done():
+			}
+		})
 	})
 }
 

pkg/app/app_test.go

@@ -105,6 +105,7 @@ func (m *mockRuntime) SetAgentThinkingLevel(context.Context, string, effort.Leve
 func (m *mockRuntime) AvailableModels(context.Context) []runtime.ModelChoice { return nil }
 func (m *mockRuntime) SupportsModelSwitching() bool                          { return false }
 func (m *mockRuntime) OnToolsChanged(func(runtime.Event))                    {}
+func (m *mockRuntime) OnBackgroundEvent(func(runtime.Event))                 {}
 
 // Verify mockRuntime implements runtime.Runtime
 var _ runtime.Runtime = (*mockRuntime)(nil)
@@ -172,6 +173,49 @@ func TestApp_Retry_SuppressesReEmittedUserMessage(t *testing.T) {
 		"only the post-StreamStarted user message should be forwarded")
 }
 
+// backgroundEventMockRuntime captures the handler App.Start registers via
+// OnBackgroundEvent so tests can emit background events through it.
+type backgroundEventMockRuntime struct {
+	mockRuntime
+
+	handler func(runtime.Event)
+}
+
+func (m *backgroundEventMockRuntime) OnBackgroundEvent(handler func(runtime.Event)) {
+	m.handler = handler
+}
+
+// TestApp_Start_ForwardsBackgroundEvents verifies Start wires the runtime's
+// out-of-band background-event hook into the app's event stream, so token
+// usage from background agent tasks reaches the TUI subscribers.
+func TestApp_Start_ForwardsBackgroundEvents(t *testing.T) {
+	t.Parallel()
+
+	rt := &backgroundEventMockRuntime{}
+	events := make(chan tea.Msg, 16)
+	app := &App{
+		runtime: rt,
+		session: session.New(),
+		events:  events,
+	}
+
+	app.Start(t.Context())
+	require.NotNil(t, rt.handler, "Start must register the background-event handler")
+
+	usage := runtime.NewTokenUsageEvent("bg-session", "worker", &runtime.Usage{
+		ContextLength: 150,
+		ContextLimit:  1000,
+	})
+	rt.handler(usage)
+
+	select {
+	case msg := <-events:
+		assert.Equal(t, usage, msg, "the background event must reach the app's event stream unchanged")
+	case <-time.After(2 * time.Second):
+		t.Fatal("timed out waiting for the forwarded background event")
+	}
+}
+
 // stubSnapshotController is a tiny SnapshotController used by the app
 // tests to drive /undo without spinning up a real shadow-git
 // repository. enabled gates SnapshotsEnabled(), and the (files, ok,

pkg/cli/runner_test.go

@@ -107,6 +107,7 @@ func (m *mockRuntime) SetAgentThinkingLevel(context.Context, string, effort.Leve
 func (m *mockRuntime) AvailableModels(context.Context) []runtime.ModelChoice                 { return nil }
 func (m *mockRuntime) SupportsModelSwitching() bool                                          { return false }
 func (m *mockRuntime) OnToolsChanged(func(runtime.Event))                                    {}
+func (m *mockRuntime) OnBackgroundEvent(func(runtime.Event))                                 {}
 func (m *mockRuntime) RegenerateTitle(context.Context, *session.Session, chan runtime.Event) {}
 
 func (m *mockRuntime) Resume(_ context.Context, req runtime.ResumeRequest) {

pkg/leantui/update_test.go

@@ -118,6 +118,7 @@ func (r *cycleThinkingRuntime) SetAgentThinkingLevel(_ context.Context, _ string
 func (r *cycleThinkingRuntime) AvailableModels(context.Context) []runtime.ModelChoice { return nil }
 func (r *cycleThinkingRuntime) SupportsModelSwitching() bool                          { return r.supports }
 func (r *cycleThinkingRuntime) OnToolsChanged(func(runtime.Event))                    {}
+func (r *cycleThinkingRuntime) OnBackgroundEvent(func(runtime.Event))                 {}
 
 var _ runtime.Runtime = (*cycleThinkingRuntime)(nil)
 

pkg/runtime/agent_delegation.go

@@ -367,6 +367,12 @@ func (r *LocalRuntime) runCollecting(ctx context.Context, parent *session.Sessio
 				onContent(choice.Content)
 			}
 		}
+		// Token usage is the one event a background sub-session surfaces
+		// out-of-band: it carries the sub-session id and agent name, so the
+		// UI can keep per-agent context accounting for background agents.
+		if usage, ok := event.(*TokenUsageEvent); ok {
+			r.emitBackgroundEvent(usage)
+		}
 		if errEvt, ok := event.(*ErrorEvent); ok {
 			errMsg = errEvt.Error
 			break

pkg/runtime/commands_test.go

@@ -92,6 +92,7 @@ func (m *mockRuntime) SetAgentThinkingLevel(context.Context, string, effort.Leve
 func (m *mockRuntime) AvailableModels(context.Context) []ModelChoice { return nil }
 func (m *mockRuntime) SupportsModelSwitching() bool                  { return false }
 func (m *mockRuntime) OnToolsChanged(func(Event))                    {}
+func (m *mockRuntime) OnBackgroundEvent(func(Event))                 {}
 
 func (m *mockRuntime) RegenerateTitle(context.Context, *session.Session, chan Event) {
 }

pkg/runtime/remote_runtime.go

@@ -702,6 +702,10 @@ func (r *RemoteRuntime) TogglePause(ctx context.Context) (bool, error) {
 // than via an out-of-band callback.
 func (r *RemoteRuntime) OnToolsChanged(func(Event)) {}
 
+// OnBackgroundEvent is a no-op for remote runtimes; background agent tasks
+// run server-side and their events are not forwarded out-of-band.
+func (r *RemoteRuntime) OnBackgroundEvent(func(Event)) {}
+
 // Close is a no-op for remote runtimes.
 func (r *RemoteRuntime) Close() error {
 	return nil