diff --git a/cli/internal/acp/server.go b/cli/internal/acp/server.go index dc944586..6e6f64ce 100644 --- a/cli/internal/acp/server.go +++ b/cli/internal/acp/server.go @@ -19,7 +19,7 @@ import ( ) type Backend interface { - Models(context.Context, string) (app.ACPModels, error) + Models(context.Context) (app.ACPModels, error) CurrentModel(context.Context, string) (string, error) SetModel(context.Context, string, string) (app.ACPModels, error) List(context.Context, *string, *string) (protocol.SessionPage, error) @@ -237,7 +237,7 @@ func (a *adapter) setup(message request) { return } // Provider discovery can be slow; do it before acquiring or pausing a lease. - models, err := a.backend.Models(a.ctx, params.SessionID) + models, err := a.backend.Models(a.ctx) if err != nil { a.respond(message.ID, nil, err) return @@ -259,7 +259,7 @@ func (a *adapter) setup(message request) { a.respond(message.ID, nil, err) return } - models.Current, err = a.backend.CurrentModel(a.ctx, attachment.SessionID) + currentModel, err := a.backend.CurrentModel(a.ctx, attachment.SessionID) if err != nil { if prior != nil { prior.resumeObservation() @@ -271,9 +271,7 @@ func (a *adapter) setup(message request) { a.respond(message.ID, nil, err) return } - if !slices.Contains(models.Available, models.Current) { - models.Available = append(models.Available, models.Current) - } + models = models.WithCurrent(currentModel) a.stopPrior(attachment.SessionID) a.mu.Lock() current := newSession(a, attachment) diff --git a/cli/internal/acp/setup_test.go b/cli/internal/acp/setup_test.go index 60a381f8..c76c4fbc 100644 --- a/cli/internal/acp/setup_test.go +++ b/cli/internal/acp/setup_test.go @@ -24,7 +24,7 @@ type setupBackend struct { calls chan string } -func (b *setupBackend) Models(context.Context, string) (app.ACPModels, error) { +func (b *setupBackend) Models(context.Context) (app.ACPModels, error) { return app.ACPModels{Current: "fixture/model", Available: []string{"fixture/model"}}, nil } @@ -154,7 +154,7 @@ type refusedResumeBackend struct { prior *session } -func (b refusedResumeBackend) Models(context.Context, string) (app.ACPModels, error) { +func (b refusedResumeBackend) Models(context.Context) (app.ACPModels, error) { return app.ACPModels{}, nil } diff --git a/cli/internal/app/acp.go b/cli/internal/app/acp.go index a5ff2497..c3368f42 100644 --- a/cli/internal/app/acp.go +++ b/cli/internal/app/acp.go @@ -19,31 +19,28 @@ type ACPModels struct { Available []string } -func (a *ACPConnection) Models(ctx context.Context, id string) (ACPModels, error) { +func (a *ACPConnection) Models(ctx context.Context) (ACPModels, error) { profiles, err := daemon.ProviderProfiles(ctx, a.connection) if err != nil { return ACPModels{}, err } - provider := profiles.Active - model := profiles.Providers[provider].Model - models := ACPModels{Current: provider + "/" + model} - if id != "" { - models.Current, err = a.CurrentModel(ctx, id) - if err != nil { - return ACPModels{}, err - } - } + models := ACPModels{} for _, catalog := range configuredModels(ctx, a.connection, profiles) { for _, model := range catalog.models { models.Available = append(models.Available, catalog.provider+"/"+model) } } - if !slices.Contains(models.Available, models.Current) { - models.Available = append(models.Available, models.Current) - } return models, nil } +func (models ACPModels) WithCurrent(current string) ACPModels { + models.Current = current + if !slices.Contains(models.Available, current) { + models.Available = append(models.Available, current) + } + return models +} + func (a *ACPConnection) CurrentModel(ctx context.Context, id string) (string, error) { session, err := daemon.GetSession(ctx, a.connection, id) return session.Provider + "/" + session.Model, err @@ -61,14 +58,15 @@ func (a *ACPConnection) SetModel(ctx context.Context, id, value string) (ACPMode if provider+"/"+model != value { return ACPModels{}, &daemon.APIError{StatusCode: 400, Code: "model_unavailable", Message: "The selected provider is no longer configured; resume to refresh models"} } - configuration, err := daemon.GetSessionConfiguration(ctx, a.connection, id) - if err != nil { + if err := switchModel(ctx, a.connection, id, modelChoice{provider, model}, a.token); err != nil { return ACPModels{}, err } - if err := daemon.SelectACPModel(ctx, a.connection, id, a.token, daemon.ModelSelectionRequest{Provider: provider, Model: model, ETag: configuration.ETag}); err != nil { + models, err := a.Models(ctx) + if err != nil { return ACPModels{}, err } - return a.Models(ctx, id) + current, err := a.CurrentModel(ctx, id) + return models.WithCurrent(current), err } func (s *Service) ConnectACP(ctx context.Context, token string) (*ACPConnection, error) { diff --git a/cli/internal/app/model_select.go b/cli/internal/app/model_select.go index 36798dc8..d7d80af3 100644 --- a/cli/internal/app/model_select.go +++ b/cli/internal/app/model_select.go @@ -77,12 +77,12 @@ func chooseModel(listed []providerModels, active, requested string) (modelChoice // switchModel moves an idle session to the chosen model without making it the // default for new sessions. -func switchModel(ctx context.Context, conn *daemon.Connection, sessionID string, choice modelChoice) error { +func switchModel(ctx context.Context, conn *daemon.Connection, sessionID string, choice modelChoice, attachmentToken string) error { configuration, err := daemon.GetSessionConfiguration(ctx, conn, sessionID) if err != nil { return err } - _, err = daemon.SelectModel(ctx, conn, sessionID, daemon.ModelSelectionRequest{Model: choice.model, Provider: choice.provider, ETag: configuration.ETag}) + _, err = daemon.SelectModel(ctx, conn, sessionID, daemon.ModelSelectionRequest{Model: choice.model, Provider: choice.provider, ETag: configuration.ETag, AttachmentToken: attachmentToken}) return err } diff --git a/cli/internal/app/prompt.go b/cli/internal/app/prompt.go index 8a103e9a..f7ffbdc3 100644 --- a/cli/internal/app/prompt.go +++ b/cli/internal/app/prompt.go @@ -56,7 +56,7 @@ func (s *Service) RunPrompt(ctx context.Context, options PromptOptions) (PromptR sessionID = session.ID } else { if model != "" { - if switchErr := switchModel(ctx, conn, sessionID, choice); switchErr != nil { + if switchErr := switchModel(ctx, conn, sessionID, choice, ""); switchErr != nil { return PromptResult{}, switchErr } } diff --git a/cli/internal/daemon/acp.go b/cli/internal/daemon/acp.go index ed534201..ef92848e 100644 --- a/cli/internal/daemon/acp.go +++ b/cli/internal/daemon/acp.go @@ -69,15 +69,6 @@ func ReleaseACP(ctx context.Context, conn *Connection, id, token string) error { }}, []int{204}, func(_ []byte, _ int) error { return nil }) } -func SelectACPModel(ctx context.Context, conn *Connection, id, token string, selection ModelSelectionRequest) error { - headers, err := observedHeaders(selection.ETag) - if err != nil { - return err - } - headers.Set("Albedo-Attachment-Token", token) - _, err = patchSessionWithHeaders(ctx, conn, id, headers, map[string]any{"model": selection.Model, "provider_profile": selection.Provider}) - return err -} func (c *ChatClient) SeedCursor(cursor protocol.Cursor) { c.mu.Lock() defer c.mu.Unlock() diff --git a/cli/internal/daemon/models.go b/cli/internal/daemon/models.go index 4c2546ec..494c1125 100644 --- a/cli/internal/daemon/models.go +++ b/cli/internal/daemon/models.go @@ -30,7 +30,7 @@ type CompactAtResult struct { Message, ETag string } type ModelSelection struct{ Provider, Model, Protocol, Effort, ETag, DefaultETag string } -type ModelSelectionRequest struct{ Model, Provider, Effort, ETag string } +type ModelSelectionRequest struct{ Model, Provider, Effort, ETag, AttachmentToken string } type ModelChangeRequest struct { Model, Provider, Effort, ETag, DefaultETag string MakeDefault bool @@ -102,7 +102,14 @@ func SelectModel(ctx context.Context, conn *Connection, id string, request Model if request.Effort != "" { body["effort"] = request.Effort } - session, err := patchSession(ctx, conn, id, request.ETag, body) + headers, err := observedHeaders(request.ETag) + if err != nil { + return ModelSelection{}, err + } + if request.AttachmentToken != "" { + headers.Set("Albedo-Attachment-Token", request.AttachmentToken) + } + session, err := patchSessionWithHeaders(ctx, conn, id, headers, body) return ModelSelection{Provider: session.Provider, Model: session.Model, Protocol: session.Protocol, Effort: session.Effort, ETag: session.ETag}, err } func ChangeModel(ctx context.Context, conn *Connection, id string, request ModelChangeRequest) (ModelSelection, error) {