feat(server): resume sessions across agent URL query changes

pkg/config/resolve.go

@@ -114,6 +114,43 @@ func resolveOne(resolvedPath string, envProvider environment.Provider) (string,
 	}
 }
 
+// StableSourceKey reduces a Sources map key to its stable identity: the part
+// that persists across variant selectors and caller/environment metadata.
+//
+// For URL references the identity is the URL's path (scheme + host + path); the
+// entire query string and fragment are treated as volatile. This is what lets a
+// session created under one variant (e.g. ?gordonTag=v9-light) resume under
+// another (?gordonTag=v9-dev) after the server is relaunched with a different
+// tag: only the query differs, so both share one identity. Treating the whole
+// query as volatile — rather than an enumerated denylist — keeps this robust as
+// new query parameters are introduced by callers.
+//
+// Keys for URL references are url.QueryEscape'd URLs; this decodes the key
+// before parsing. For keys that are not URL references (local files, OCI refs,
+// built-ins) or that cannot be decoded/parsed, the key is returned unchanged,
+// so callers can compare identities without special-casing the source type.
+//
+// The resume fallback that consumes this (see the session manager's
+// resolveSource) always prefers an exact key match and only uses the stable
+// identity when it resolves unambiguously, so collapsing distinct query strings
+// to one identity never silently selects the wrong side-by-side variant.
+func StableSourceKey(key string) string {
+	decoded, err := url.QueryUnescape(key)
+	if err != nil {
+		return key
+	}
+	if !IsURLReference(decoded) {
+		return key
+	}
+	u, err := url.Parse(decoded)
+	if err != nil {
+		return key
+	}
+	u.RawQuery = ""
+	u.Fragment = ""
+	return u.String()
+}
+
 // resolveDirectory enumerates YAML files in a directory and resolves each one.
 func resolveDirectory(dirPath string, envProvider environment.Provider) (Sources, error) {
 	entries, err := os.ReadDir(dirPath)

pkg/config/resolve_test.go

@@ -805,3 +805,69 @@ func TestResolveSources_URLEncodedKey(t *testing.T) {
 	// The source Name() should still return the original URL for fetching
 	assert.Equal(t, testURL, source.Name())
 }
+
+func TestStableSourceKey(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name string
+		key  string
+		want string
+	}{
+		{
+			name: "strips query from url key",
+			key:  url.QueryEscape("http://localhost:7777/gordon-agent?gordonTag=v9-light&desktopVersion=4.81.0&origin=desktop"),
+			want: "http://localhost:7777/gordon-agent",
+		},
+		{
+			name: "another variant normalises to the same identity",
+			key:  url.QueryEscape("http://localhost:7777/gordon-agent?gordonTag=v9-dev&desktopVersion=4.81.0&origin=desktop"),
+			want: "http://localhost:7777/gordon-agent",
+		},
+		{
+			name: "strips all query params, keeping the path identity",
+			key:  url.QueryEscape("http://localhost:7777/gordon-agent?team=blue&gordonTag=v9"),
+			want: "http://localhost:7777/gordon-agent",
+		},
+		{
+			name: "strips fragment as well",
+			key:  url.QueryEscape("http://localhost:7777/gordon-agent?gordonTag=v9#section"),
+			want: "http://localhost:7777/gordon-agent",
+		},
+		{
+			name: "distinct paths keep distinct identities",
+			key:  url.QueryEscape("http://localhost:7777/other-agent?gordonTag=v9"),
+			want: "http://localhost:7777/other-agent",
+		},
+		{
+			name: "non-url key is returned unchanged",
+			key:  "docker_gordon.yaml",
+			want: "docker_gordon.yaml",
+		},
+		{
+			name: "local file key is returned unchanged",
+			key:  "my-agent",
+			want: "my-agent",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			assert.Equal(t, tt.want, StableSourceKey(tt.key))
+		})
+	}
+}
+
+// TestStableSourceKey_VariantsCollide is the property the resume fallback
+// relies on: two source keys that differ only by volatile query parameters
+// must produce the same stable identity.
+func TestStableSourceKey_VariantsCollide(t *testing.T) {
+	t.Parallel()
+
+	light := url.QueryEscape("http://localhost:7777/gordon-agent?gordonTag=v9-light&origin=desktop")
+	dev := url.QueryEscape("http://localhost:7777/gordon-agent?gordonTag=v9-dev&origin=desktop")
+
+	assert.Equal(t, StableSourceKey(light), StableSourceKey(dev))
+}

pkg/server/resolve_source_test.go

@@ -0,0 +1,108 @@
+package server
+
+import (
+	"net/url"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/docker/docker-agent/pkg/config"
+)
+
+func gordonKey(tag string) string {
+	return url.QueryEscape("http://localhost:7777/gordon-agent?gordonTag=" + tag + "&desktopVersion=4.81.0&origin=desktop")
+}
+
+// TestResolveSource_ExactMatchPreferred verifies that an exact key match is
+// always used, even when other variants exist that would also normalise to the
+// same stable identity.
+func TestResolveSource_ExactMatchPreferred(t *testing.T) {
+	t.Parallel()
+
+	light := config.NewBytesSource("light", []byte("light"))
+	dev := config.NewBytesSource("dev", []byte("dev"))
+	sm := &SessionManager{
+		Sources: config.Sources{
+			gordonKey("v9-light"): light,
+			gordonKey("v9-dev"):   dev,
+		},
+	}
+
+	got, err := sm.resolveSource(gordonKey("v9-dev"))
+	require.NoError(t, err)
+	assert.Equal(t, "dev", got.Name())
+}
+
+// TestResolveSource_FallbackAcrossTag is the resume-after-upgrade case: the
+// session recorded the v9-light key, but the server was relaunched with only
+// v9-dev. The fallback resolves it because the two share a stable identity.
+func TestResolveSource_FallbackAcrossTag(t *testing.T) {
+	t.Parallel()
+
+	dev := config.NewBytesSource("dev", []byte("dev"))
+	sm := &SessionManager{
+		Sources: config.Sources{
+			gordonKey("v9-dev"): dev,
+		},
+	}
+
+	got, err := sm.resolveSource(gordonKey("v9-light"))
+	require.NoError(t, err)
+	assert.Equal(t, "dev", got.Name(), "an old session's tagged ref should resolve to the live source")
+}
+
+// TestResolveSource_AmbiguousFallbackFails verifies that the fallback refuses
+// to guess: when several live sources share the requested stable identity and
+// none matches exactly, it returns not-found rather than picking one.
+func TestResolveSource_AmbiguousFallbackFails(t *testing.T) {
+	t.Parallel()
+
+	sm := &SessionManager{
+		Sources: config.Sources{
+			gordonKey("v9-light"): config.NewBytesSource("light", []byte("light")),
+			gordonKey("v9-dev"):   config.NewBytesSource("dev", []byte("dev")),
+		},
+	}
+
+	// A third tag not present exactly; two candidates share its identity.
+	_, err := sm.resolveSource(gordonKey("v9-canary"))
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "agent not found")
+}
+
+// TestResolveSource_NoMatchFails verifies the plain not-found path is preserved
+// for references that share no stable identity with any live source.
+func TestResolveSource_NoMatchFails(t *testing.T) {
+	t.Parallel()
+
+	sm := &SessionManager{
+		Sources: config.Sources{
+			gordonKey("v9-dev"): config.NewBytesSource("dev", []byte("dev")),
+		},
+	}
+
+	_, err := sm.resolveSource(url.QueryEscape("http://localhost:7777/other-agent?gordonTag=v9-dev"))
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "agent not found")
+}
+
+// TestResolveSource_LocalFileKeyExactOnly verifies that non-URL keys (local
+// files, OCI refs) still resolve only by exact match, since their stable
+// identity is the key itself.
+func TestResolveSource_LocalFileKeyExactOnly(t *testing.T) {
+	t.Parallel()
+
+	sm := &SessionManager{
+		Sources: config.Sources{
+			"my-agent": config.NewBytesSource("my-agent", []byte("x")),
+		},
+	}
+
+	got, err := sm.resolveSource("my-agent")
+	require.NoError(t, err)
+	assert.Equal(t, "my-agent", got.Name())
+
+	_, err = sm.resolveSource("other-agent")
+	require.Error(t, err)
+}

pkg/server/session_manager.go

@@ -1000,9 +1000,9 @@ func (sm *SessionManager) runtimeForSession(ctx context.Context, sess *session.S
 }
 
 func (sm *SessionManager) loadTeam(ctx context.Context, agentFilename string, runConfig *config.RuntimeConfig) (*team.Team, error) {
-	agentSource, found := sm.Sources[agentFilename]
-	if !found {
-		return nil, fmt.Errorf("agent not found: %s", agentFilename)
+	agentSource, err := sm.resolveSource(agentFilename)
+	if err != nil {
+		return nil, err
 	}
 
 	return teamloader.Load(ctx, agentSource, runConfig, loaderdefaults.Opts()...)
@@ -1011,14 +1011,48 @@ func (sm *SessionManager) loadTeam(ctx context.Context, agentFilename string, ru
 // loadTeamWithConfig is like loadTeam but also returns the loaded model and
 // provider configuration so the runtime can be wired for model switching.
 func (sm *SessionManager) loadTeamWithConfig(ctx context.Context, agentFilename string, runConfig *config.RuntimeConfig) (*teamloader.LoadResult, error) {
-	agentSource, found := sm.Sources[agentFilename]
-	if !found {
-		return nil, fmt.Errorf("agent not found: %s", agentFilename)
+	agentSource, err := sm.resolveSource(agentFilename)
+	if err != nil {
+		return nil, err
 	}
 
 	return teamloader.LoadWithConfig(ctx, agentSource, runConfig, loaderdefaults.Opts()...)
 }
 
+// resolveSource looks up the agent source for agentFilename.
+//
+// An exact match is always preferred so that distinct variants served side by
+// side (e.g. two gordonTag values in the same process) keep their own sources.
+// When there is no exact match, it falls back to matching on a stable identity
+// that ignores volatile URL query parameters (see config.StableSourceKey). This
+// lets a session created under one variant resume under another after the
+// server is relaunched with a different tag — the exact key recorded by the
+// client no longer exists, but the underlying agent does.
+//
+// The fallback only fires when it is unambiguous: if several live sources share
+// the same stable identity, resolving would be a guess, so it returns the
+// not-found error instead of silently picking one.
+func (sm *SessionManager) resolveSource(agentFilename string) (config.Source, error) {
+	if agentSource, found := sm.Sources[agentFilename]; found {
+		return agentSource, nil
+	}
+
+	want := config.StableSourceKey(agentFilename)
+	var match config.Source
+	var matches int
+	for key, source := range sm.Sources {
+		if config.StableSourceKey(key) == want {
+			match = source
+			matches++
+		}
+	}
+	if matches == 1 {
+		return match, nil
+	}
+
+	return nil, fmt.Errorf("agent not found: %s", agentFilename)
+}
+
 // applyRunModelOverride applies modelRef as the per-agent model override
 // on the session backing rs. It mirrors the in-memory mutations that
 // SetSessionAgentModel performs, but without acquiring sm.mux (the