Skip to content

Commit b711608

Browse files
committed
feat(tui): add argument auto-completion for /toolset-restart
Pressing Tab after "/toolset-restart " opens a completion popup listing the current agent's toolsets. Non-restartable toolsets are shown but dimmed and annotated "(not restartable)", and cannot be selected. The completion opens only on Tab, never auto-opening on space. The argument-completion hook (ArgumentCandidate/CompleteArgument on commands.Item plus the ArgumentCompleter source interface) is command-agnostic so it can be reused for other commands (e.g. /effort, /drop) later, though only /toolset-restart is wired up here.
1 parent 66a9488 commit b711608

18 files changed

Lines changed: 1134 additions & 18 deletions

pkg/runtime/runtime.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1139,6 +1139,7 @@ func toolsetStatusFor(ts tools.ToolSet) tools.ToolsetStatus {
11391139
// earlier if Start failed.
11401140
status.State = lifecycleStateForUnsupervised(ts)
11411141
}
1142+
_, status.Restartable = tools.As[tools.Restartable](ts)
11421143
status.Name = nameFor(ts, status.Description)
11431144
return status
11441145
}

pkg/runtime/toolsetstatus_test.go

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,23 @@ func TestToolsetStatusFor_NoStatableMeansReady(t *testing.T) {
5959
assert.Equal(t, 0, got.RestartCount)
6060
require.NoError(t, got.LastError)
6161
assert.Equal(t, "filesystem", got.Description)
62+
assert.False(t, got.Restartable, "toolset without Restart() must report Restartable=false")
63+
}
64+
65+
func TestToolsetStatusFor_RestartableReportsTrue(t *testing.T) {
66+
t.Parallel()
67+
68+
ts := &restartableToolset{desc: "mcp(stdio cmd=foo)", state: lifecycle.StateInfo{State: lifecycle.StateReady}}
69+
got := toolsetStatusFor(ts)
70+
assert.True(t, got.Restartable)
71+
}
72+
73+
func TestToolsetStatusFor_NonRestartableReportsFalse(t *testing.T) {
74+
t.Parallel()
75+
76+
ts := &statefulToolset{desc: "filesystem", info: lifecycle.StateInfo{State: lifecycle.StateReady}}
77+
got := toolsetStatusFor(ts)
78+
assert.False(t, got.Restartable)
6279
}
6380

6481
// TestToolsetStatusFor_UnwrapsStartable verifies the inner Statable is

pkg/tools/status.go

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,4 +43,9 @@ type ToolsetStatus struct {
4343
// RestartCount is the number of supervisor restarts since the last
4444
// successful Ready transition. Zero for toolsets without a supervisor.
4545
RestartCount int
46+
// Restartable reports whether the toolset implements Restartable
47+
// (typically supervisor-backed MCP/LSP toolsets). Surfaces such as the
48+
// /toolset-restart completion popup use this to annotate, not hide,
49+
// entries that cannot actually be restarted.
50+
Restartable bool
4651
}

pkg/tui/commands/commands.go

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package commands
22

33
import (
4+
"cmp"
45
"context"
56
"fmt"
67
"slices"
@@ -11,6 +12,7 @@ import (
1112
"github.com/docker/docker-agent/pkg/app"
1213
"github.com/docker/docker-agent/pkg/config/types"
1314
"github.com/docker/docker-agent/pkg/feedback"
15+
"github.com/docker/docker-agent/pkg/tools"
1416
mcptools "github.com/docker/docker-agent/pkg/tools/mcp"
1517
"github.com/docker/docker-agent/pkg/tui/components/toolcommon"
1618
"github.com/docker/docker-agent/pkg/tui/core"
@@ -26,6 +28,18 @@ type Category struct {
2628
Commands []Item
2729
}
2830

31+
// ArgumentCandidate is one completable value for a command's argument,
32+
// e.g. a toolset name for /toolset-restart. It is intentionally free of any
33+
// completion-UI or app-domain types so pkg/tui/commands stays decoupled from
34+
// both pkg/tui/components/completion and pkg/app.
35+
type ArgumentCandidate struct {
36+
Label string
37+
Description string
38+
// Disabled marks a candidate that is shown for context but cannot be
39+
// submitted (e.g. a non-restartable toolset for /toolset-restart).
40+
Disabled bool
41+
}
42+
2943
// Item represents a single command in the palette
3044
type Item struct {
3145
ID string
@@ -38,6 +52,11 @@ type Item struct {
3852
// Immediate marks commands that should run as soon as they are submitted
3953
// instead of being treated as ordinary queued chat input.
4054
Immediate bool
55+
// CompleteArgument, when set, returns the candidates for this command's
56+
// argument. Called lazily at completion-popup-open time so results
57+
// reflect current runtime state (e.g. toolset lifecycle). Nil for
58+
// commands with no argument completion.
59+
CompleteArgument func() []ArgumentCandidate
4160
}
4261

4362
func builtInSessionCommands() []Item {
@@ -577,13 +596,72 @@ func newMCPPromptItem(promptName string, promptInfo mcptools.PromptInfo) Item {
577596
}
578597
}
579598

599+
// toolsetStatusSource is the minimal surface toolsetRestartCandidates needs.
600+
// *app.App satisfies it; tests can supply a stub instead of constructing a
601+
// full application.
602+
type toolsetStatusSource interface {
603+
CurrentAgentToolsetStatuses() []tools.ToolsetStatus
604+
}
605+
606+
// toolsetRestartCandidates returns one ArgumentCandidate per toolset of the
607+
// current agent, in declaration order, deduplicated by name (preferring the
608+
// restartable entry when duplicate names disagree — display names are not
609+
// guaranteed unique, mirroring runtime.RestartToolset's own matching).
610+
// Non-restartable toolsets are marked Disabled rather than omitted, so the
611+
// popup teaches users the full toolset inventory and why a restart isn't
612+
// offered for some of them.
613+
func toolsetRestartCandidates(source toolsetStatusSource) []ArgumentCandidate {
614+
statuses := source.CurrentAgentToolsetStatuses()
615+
byName := make(map[string]tools.ToolsetStatus, len(statuses))
616+
order := make([]string, 0, len(statuses))
617+
for _, s := range statuses {
618+
if s.Name == "" {
619+
continue
620+
}
621+
if existing, ok := byName[s.Name]; !ok {
622+
byName[s.Name] = s
623+
order = append(order, s.Name)
624+
} else if !existing.Restartable && s.Restartable {
625+
byName[s.Name] = s
626+
}
627+
}
628+
629+
candidates := make([]ArgumentCandidate, 0, len(order))
630+
for _, name := range order {
631+
s := byName[name]
632+
candidates = append(candidates, ArgumentCandidate{
633+
Label: s.Name,
634+
Description: cmp.Or(s.Kind, "Built-in") + " · " + s.State.String(),
635+
Disabled: !s.Restartable,
636+
})
637+
}
638+
return candidates
639+
}
640+
641+
// attachToolsetRestartCompletion wires the toolset-name argument completer
642+
// onto the /toolset-restart item. Attached post-hoc (rather than inline in
643+
// builtInSessionCommands) so that function stays free of any status-source
644+
// dependency.
645+
func attachToolsetRestartCompletion(items []Item, source toolsetStatusSource) {
646+
for i := range items {
647+
if items[i].ID != "session.toolset.restart" {
648+
continue
649+
}
650+
items[i].CompleteArgument = func() []ArgumentCandidate {
651+
return toolsetRestartCandidates(source)
652+
}
653+
return
654+
}
655+
}
656+
580657
// BuildCommandCategories builds the list of command categories for the command palette
581658
func BuildCommandCategories(ctx context.Context, application *app.App) []Category {
582659
// Get session commands and filter based on model capabilities
583660
sessionCommands := builtInSessionCommands()
584661
if !application.SnapshotsEnabled() {
585662
sessionCommands = removeByIDs(sessionCommands, snapshotCommandIDs)
586663
}
664+
attachToolsetRestartCompletion(sessionCommands, application)
587665

588666
categories := []Category{
589667
{

pkg/tui/commands/commands_test.go

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@ import (
66
"github.com/stretchr/testify/assert"
77
"github.com/stretchr/testify/require"
88

9+
"github.com/docker/docker-agent/pkg/tools"
10+
"github.com/docker/docker-agent/pkg/tools/lifecycle"
911
"github.com/docker/docker-agent/pkg/tui/messages"
1012
)
1113

@@ -277,3 +279,122 @@ func TestSettingsCommandLabel(t *testing.T) {
277279
assert.Equal(t, "settings.open", settings.ID)
278280
assert.Equal(t, "Settings", settings.Label)
279281
}
282+
283+
// stubToolsetStatusSource is a minimal toolsetStatusSource for exercising
284+
// toolsetRestartCandidates without a real *app.App.
285+
type stubToolsetStatusSource struct {
286+
statuses []tools.ToolsetStatus
287+
}
288+
289+
func (s stubToolsetStatusSource) CurrentAgentToolsetStatuses() []tools.ToolsetStatus {
290+
return s.statuses
291+
}
292+
293+
func TestToolsetRestartCandidates_AllShownWithDisabledFlag(t *testing.T) {
294+
t.Parallel()
295+
296+
source := stubToolsetStatusSource{statuses: []tools.ToolsetStatus{
297+
{Name: "github", Kind: "MCP", State: lifecycle.StateReady, Restartable: true},
298+
{Name: "filesystem", State: lifecycle.StateReady, Restartable: false},
299+
}}
300+
301+
candidates := toolsetRestartCandidates(source)
302+
require.Len(t, candidates, 2)
303+
304+
assert.Equal(t, "github", candidates[0].Label)
305+
assert.False(t, candidates[0].Disabled)
306+
assert.Contains(t, candidates[0].Description, "MCP")
307+
308+
assert.Equal(t, "filesystem", candidates[1].Label)
309+
assert.True(t, candidates[1].Disabled, "non-restartable toolsets must be shown, not hidden")
310+
assert.Contains(t, candidates[1].Description, "Built-in")
311+
}
312+
313+
func TestToolsetRestartCandidates_DedupePrefersRestartable(t *testing.T) {
314+
t.Parallel()
315+
316+
source := stubToolsetStatusSource{statuses: []tools.ToolsetStatus{
317+
{Name: "dup", State: lifecycle.StateStopped, Restartable: false},
318+
{Name: "dup", State: lifecycle.StateReady, Restartable: true},
319+
}}
320+
321+
candidates := toolsetRestartCandidates(source)
322+
require.Len(t, candidates, 1, "duplicate names must be deduplicated")
323+
assert.False(t, candidates[0].Disabled, "the restartable entry must win over the non-restartable duplicate")
324+
}
325+
326+
func TestToolsetRestartCandidates_EmptyNameSkipped(t *testing.T) {
327+
t.Parallel()
328+
329+
source := stubToolsetStatusSource{statuses: []tools.ToolsetStatus{{Name: "", Restartable: true}}}
330+
assert.Empty(t, toolsetRestartCandidates(source))
331+
}
332+
333+
func TestAttachToolsetRestartCompletion(t *testing.T) {
334+
t.Parallel()
335+
336+
items := builtInSessionCommands()
337+
source := stubToolsetStatusSource{statuses: []tools.ToolsetStatus{
338+
{Name: "github", Restartable: true},
339+
}}
340+
attachToolsetRestartCompletion(items, source)
341+
342+
var restart, other *Item
343+
for i := range items {
344+
switch items[i].ID {
345+
case "session.toolset.restart":
346+
restart = &items[i]
347+
case "session.exit":
348+
other = &items[i]
349+
}
350+
}
351+
require.NotNil(t, restart)
352+
require.NotNil(t, restart.CompleteArgument, "the restart item must get a completer")
353+
candidates := restart.CompleteArgument()
354+
require.Len(t, candidates, 1)
355+
assert.Equal(t, "github", candidates[0].Label)
356+
357+
require.NotNil(t, other)
358+
assert.Nil(t, other.CompleteArgument, "commands without a provider keep a nil CompleteArgument")
359+
}
360+
361+
// TestAttachToolsetRestartCompletion_QueriesFreshOnEachCall is a regression
362+
// test for bug 2 (PR #3728): the attached CompleteArgument closure must query
363+
// the toolset status source live on every call, not capture a snapshot at
364+
// attach time. Otherwise the first Tab after "/toolset-restart " shows
365+
// whatever toolsets existed when the command list was built, not the
366+
// current ones.
367+
func TestAttachToolsetRestartCompletion_QueriesFreshOnEachCall(t *testing.T) {
368+
t.Parallel()
369+
370+
items := builtInSessionCommands()
371+
source := &stubToolsetStatusSource{statuses: []tools.ToolsetStatus{
372+
{Name: "github", State: lifecycle.StateReady, Restartable: true},
373+
}}
374+
attachToolsetRestartCompletion(items, source)
375+
376+
var restart *Item
377+
for i := range items {
378+
if items[i].ID == "session.toolset.restart" {
379+
restart = &items[i]
380+
}
381+
}
382+
require.NotNil(t, restart)
383+
384+
first := restart.CompleteArgument()
385+
require.Len(t, first, 1)
386+
assert.Equal(t, "github", first[0].Label)
387+
388+
// The toolset status source changes after the completer was attached
389+
// (e.g. the agent restarted with a different toolset set). A second call
390+
// must reflect that change rather than replaying the first snapshot.
391+
source.statuses = []tools.ToolsetStatus{
392+
{Name: "filesystem", State: lifecycle.StateReady, Restartable: true},
393+
{Name: "github", State: lifecycle.StateReady, Restartable: true},
394+
}
395+
396+
second := restart.CompleteArgument()
397+
require.Len(t, second, 2, "must reflect the current toolset set, not the one captured at attach time")
398+
assert.Equal(t, "filesystem", second[0].Label)
399+
assert.Equal(t, "github", second[1].Label)
400+
}

0 commit comments

Comments
 (0)