Skip to content

Commit ddfdf28

Browse files
PierrunoYTampagent
andauthored
Fix ChatGPT OAuth model discovery and ACP model selection (#724)
* fix: expose authenticated ChatGPT models to ACP Co-authored-by: Pierre Bruno <pierrebruno@hotmail.ch> * fix(acp): persist discovered model selections * fix(acp): normalize model selections --------- Co-authored-by: Amp <amp@ampcode.com>
1 parent 96859c9 commit ddfdf28

11 files changed

Lines changed: 592 additions & 71 deletions

File tree

‎internal/acp/agent.go‎

Lines changed: 178 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,9 @@ import (
1111

1212
"github.com/Gitlawb/zero/internal/agent"
1313
"github.com/Gitlawb/zero/internal/config"
14+
"github.com/Gitlawb/zero/internal/providercatalog"
15+
"github.com/Gitlawb/zero/internal/providermodelcatalog"
16+
"github.com/Gitlawb/zero/internal/providermodeldiscovery"
1417
"github.com/Gitlawb/zero/internal/sandbox"
1518
"github.com/Gitlawb/zero/internal/sessions"
1619
"github.com/Gitlawb/zero/internal/tools"
@@ -23,9 +26,10 @@ import (
2326
// means the editor only hosts the thread — ZERO owns BYOK and telemetry-free
2427
// operation.
2528
type Deps struct {
26-
ResolveConfig func(workspaceRoot string, overrides config.Overrides) (config.ResolvedConfig, error)
27-
NewProvider func(profile config.ProviderProfile) (zeroruntime.Provider, error)
28-
RunAgent func(ctx context.Context, prompt string, provider zeroruntime.Provider, opts agent.Options) (agent.Result, error)
29+
ResolveConfig func(workspaceRoot string, overrides config.Overrides) (config.ResolvedConfig, error)
30+
DiscoverModels func(context.Context, config.ProviderProfile) ([]providermodeldiscovery.Model, error)
31+
NewProvider func(profile config.ProviderProfile) (zeroruntime.Provider, error)
32+
RunAgent func(ctx context.Context, prompt string, provider zeroruntime.Provider, opts agent.Options) (agent.Result, error)
2933
// BuildWorkspace builds the SCOPED tool registry and the sandbox engine for a
3034
// validated workspace root, so ACP shell tools (bash/exec_command) are confined
3135
// exactly like the exec surface — never run unconfined on the host.
@@ -60,13 +64,16 @@ type acpSession struct {
6064
// turnMu serializes prompt turns for one session: concurrent session/prompt
6165
// calls run one at a time so they can't interleave history or clobber the
6266
// single cancel slot.
63-
turnMu sync.Mutex
67+
turnMu sync.Mutex
68+
modelMu sync.Mutex
6469

65-
mu sync.Mutex
66-
mode agent.PermissionMode
67-
model string // override; "" => config default
68-
cancel context.CancelFunc
69-
history []turnRecord
70+
mu sync.Mutex
71+
mode agent.PermissionMode
72+
model string
73+
models []SessionConfigOptionValue
74+
restrictModels bool
75+
cancel context.CancelFunc
76+
history []turnRecord
7077
}
7178

7279
// NewAgent builds the ACP server and registers its method handlers on conn.
@@ -120,7 +127,7 @@ func (a *Agent) handleInitialize(_ context.Context, params json.RawMessage) (any
120127

121128
// ---- session lifecycle ----
122129

123-
func (a *Agent) handleSessionNew(_ context.Context, params json.RawMessage) (any, error) {
130+
func (a *Agent) handleSessionNew(ctx context.Context, params json.RawMessage) (any, error) {
124131
var p NewSessionParams
125132
if err := json.Unmarshal(params, &p); err != nil {
126133
return nil, RPCError(codeInvalidParams, "invalid session/new params")
@@ -129,18 +136,23 @@ func (a *Agent) handleSessionNew(_ context.Context, params json.RawMessage) (any
129136
if err != nil {
130137
return nil, RPCError(codeInvalidParams, err.Error())
131138
}
132-
meta, err := a.deps.Store.Create(sessions.CreateInput{Title: "ACP session", Cwd: root})
139+
model, models, restrictModels, err := a.resolveModelChoices(ctx, root)
140+
if err != nil {
141+
return nil, RPCError(codeInternalError, "config: "+err.Error())
142+
}
143+
meta, err := a.deps.Store.Create(sessions.CreateInput{Title: "ACP session", Cwd: root, ModelID: model})
133144
if err != nil {
134145
return nil, RPCError(codeInternalError, "create session: "+err.Error())
135146
}
136-
sess := a.registerSession(meta.SessionID, root, nil)
147+
sess := a.registerSession(meta.SessionID, root, nil, model, models, restrictModels)
137148
return NewSessionResult{
138-
SessionID: sess.id,
139-
Modes: a.modeState(sess),
149+
SessionID: sess.id,
150+
ConfigOptions: a.configOptions(sess),
151+
Modes: a.modeState(sess),
140152
}, nil
141153
}
142154

143-
func (a *Agent) handleSessionLoad(_ context.Context, params json.RawMessage) (any, error) {
155+
func (a *Agent) handleSessionLoad(ctx context.Context, params json.RawMessage) (any, error) {
144156
var p LoadSessionParams
145157
if err := json.Unmarshal(params, &p); err != nil {
146158
return nil, RPCError(codeInvalidParams, "invalid session/load params")
@@ -161,15 +173,26 @@ func (a *Agent) handleSessionLoad(_ context.Context, params json.RawMessage) (an
161173
// a half-initialized session (registerSession sets history under the lock and
162174
// reuses an already-live session rather than orphaning its in-flight turn).
163175
history, historyErr := a.loadHistory(meta.SessionID)
164-
sess := a.registerSession(meta.SessionID, root, history)
176+
model, models, restrictModels, err := a.resolveModelChoices(ctx, root)
177+
if err != nil {
178+
return nil, RPCError(codeInternalError, "config: "+err.Error())
179+
}
180+
if persistedModel := strings.TrimSpace(meta.ModelID); persistedModel != "" && (!restrictModels || modelChoiceExists(models, persistedModel)) {
181+
model = persistedModel
182+
if !modelChoiceExists(models, persistedModel) {
183+
models = append(models, SessionConfigOptionValue{Value: persistedModel, Name: persistedModel})
184+
}
185+
}
186+
sess := a.registerSession(meta.SessionID, root, history, model, models, restrictModels)
165187
a.warnPersistence(
166188
&notifier{conn: a.conn, sessionID: sess.id},
167189
"load session history",
168190
"Could not load session history. The session is open, but earlier turns may be missing until storage recovers.",
169191
historyErr,
170192
)
171193
return LoadSessionResult{
172-
Modes: a.modeState(sess),
194+
ConfigOptions: a.configOptions(sess),
195+
Modes: a.modeState(sess),
173196
}, nil
174197
}
175198

@@ -354,11 +377,27 @@ func (a *Agent) handleSetConfigOption(_ context.Context, params json.RawMessage)
354377
if sess == nil {
355378
return nil, RPCError(codeInvalidParams, "unknown session: "+p.SessionID)
356379
}
357-
if p.ConfigID != configIDModel {
380+
switch p.ConfigID {
381+
case configIDModel:
382+
model := strings.TrimSpace(p.Value)
383+
if err := a.updateModel(sess, model, sess.restrictModels); err != nil {
384+
return nil, err
385+
}
386+
case configIDMode:
387+
mode := agent.PermissionMode(p.Value)
388+
switch mode {
389+
case agent.PermissionModeAuto, agent.PermissionModeAsk:
390+
sess.setMode(mode)
391+
(&notifier{conn: a.conn, sessionID: sess.id}).currentMode(string(mode))
392+
case agent.PermissionModeUnsafe:
393+
return nil, RPCError(codeInvalidParams, "mode not permitted over ACP: "+p.Value)
394+
default:
395+
return nil, RPCError(codeInvalidParams, "unknown mode: "+p.Value)
396+
}
397+
default:
358398
return nil, RPCError(codeInvalidParams, "unknown config option: "+p.ConfigID)
359399
}
360-
sess.setModel(p.Value)
361-
return SetSessionConfigOptionResult{}, nil
400+
return SetSessionConfigOptionResult{ConfigOptions: a.configOptions(sess)}, nil
362401
}
363402

364403
func (a *Agent) handleZeroSetModel(_ context.Context, params json.RawMessage) (any, error) {
@@ -370,8 +409,24 @@ func (a *Agent) handleZeroSetModel(_ context.Context, params json.RawMessage) (a
370409
if sess == nil {
371410
return nil, RPCError(codeInvalidParams, "unknown session: "+p.SessionID)
372411
}
373-
sess.setModel(p.Model)
374-
return ZeroSetModelResult{Model: p.Model}, nil
412+
model := strings.TrimSpace(p.Model)
413+
if err := a.updateModel(sess, model, false); err != nil {
414+
return nil, err
415+
}
416+
return ZeroSetModelResult{Model: model}, nil
417+
}
418+
419+
func (a *Agent) updateModel(sess *acpSession, model string, restrictModels bool) error {
420+
sess.modelMu.Lock()
421+
defer sess.modelMu.Unlock()
422+
if restrictModels && !sess.hasModel(model) {
423+
return RPCError(codeInvalidParams, "unknown model: "+model)
424+
}
425+
if _, err := a.deps.Store.UpdateModel(sess.id, model); err != nil {
426+
return RPCError(codeInternalError, "save model selection: "+err.Error())
427+
}
428+
sess.setModel(model)
429+
return nil
375430
}
376431

377432
func (a *Agent) handleCancel(_ context.Context, params json.RawMessage) {
@@ -398,6 +453,76 @@ func (a *Agent) modeState(s *acpSession) *SessionModeState {
398453
}
399454
}
400455

456+
// resolveModelChoices advertises authenticated live models when discovery is
457+
// available. The configured model is always retained as the sole fallback.
458+
func (a *Agent) resolveModelChoices(ctx context.Context, cwd string) (string, []SessionConfigOptionValue, bool, error) {
459+
resolved, err := a.deps.ResolveConfig(cwd, config.Overrides{})
460+
if err != nil {
461+
return "", nil, false, err
462+
}
463+
selected := strings.TrimSpace(resolved.Provider.Model)
464+
options := make([]SessionConfigOptionValue, 0, 8)
465+
seen := make(map[string]bool)
466+
add := func(id, description string) {
467+
id = strings.TrimSpace(id)
468+
if id == "" || seen[id] {
469+
return
470+
}
471+
seen[id] = true
472+
options = append(options, SessionConfigOptionValue{Value: id, Name: id, Description: description})
473+
}
474+
add(selected, "")
475+
descriptor, knownProvider := providercatalog.Get(resolved.Provider.CatalogID)
476+
restrictModels := knownProvider && !descriptor.Custom
477+
if a.deps.DiscoverModels != nil {
478+
discovered, discoverErr := a.deps.DiscoverModels(ctx, resolved.Provider)
479+
if ctx.Err() != nil {
480+
return "", nil, false, ctx.Err()
481+
}
482+
if discoverErr == nil && len(discovered) > 0 {
483+
for _, model := range discovered {
484+
if providermodelcatalog.ModelIDAllowedForProvider(resolved.Provider.CatalogID, model.ID) {
485+
add(model.ID, model.Description)
486+
}
487+
}
488+
}
489+
}
490+
return selected, options, restrictModels, nil
491+
}
492+
493+
func modelChoiceExists(models []SessionConfigOptionValue, model string) bool {
494+
for _, option := range models {
495+
if option.Value == model {
496+
return true
497+
}
498+
}
499+
return false
500+
}
501+
502+
func (a *Agent) configOptions(s *acpSession) []SessionConfigOption {
503+
model, models, mode := s.configState()
504+
return []SessionConfigOption{{
505+
ID: configIDModel,
506+
Name: "Model",
507+
Description: "Model used for this session.",
508+
Category: configCategoryModel,
509+
Type: configOptionTypeSelect,
510+
CurrentValue: model,
511+
Options: models,
512+
}, {
513+
ID: configIDMode,
514+
Name: "Mode",
515+
Description: "Permission mode used for this session.",
516+
Category: configCategoryMode,
517+
Type: configOptionTypeSelect,
518+
CurrentValue: string(mode),
519+
Options: []SessionConfigOptionValue{
520+
{Value: string(agent.PermissionModeAuto), Name: "Auto", Description: "Run safe tools automatically; ask before risky ones."},
521+
{Value: string(agent.PermissionModeAsk), Name: "Ask", Description: "Ask before every tool that changes state."},
522+
},
523+
}}
524+
}
525+
401526
// ---- persistence + continuity ----
402527

403528
func (a *Agent) persistTurn(sess *acpSession, user, assistant string) error {
@@ -525,13 +650,13 @@ func promptImages(blocks []ContentBlock) []zeroruntime.ImageBlock {
525650
// session is returned unchanged rather than orphaning its turn or resetting its
526651
// mode/model. history is set BEFORE publishing so no concurrent prompt can read a
527652
// half-initialized session.
528-
func (a *Agent) registerSession(id, cwd string, history []turnRecord) *acpSession {
653+
func (a *Agent) registerSession(id, cwd string, history []turnRecord, model string, models []SessionConfigOptionValue, restrictModels bool) *acpSession {
529654
a.mu.Lock()
530655
defer a.mu.Unlock()
531656
if existing := a.sessions[id]; existing != nil {
532657
return existing
533658
}
534-
sess := &acpSession{id: id, cwd: cwd, mode: agent.PermissionModeAuto, history: history}
659+
sess := &acpSession{id: id, cwd: cwd, mode: agent.PermissionModeAuto, model: model, models: models, restrictModels: restrictModels, history: history}
535660
a.sessions[id] = sess
536661
return sess
537662
}
@@ -570,7 +695,18 @@ func (s *acpSession) currentMode() agent.PermissionMode {
570695
}
571696

572697
func (s *acpSession) setModel(model string) {
698+
model = strings.TrimSpace(model)
573699
s.mu.Lock()
700+
found := false
701+
for _, option := range s.models {
702+
if option.Value == model {
703+
found = true
704+
break
705+
}
706+
}
707+
if !found && model != "" {
708+
s.models = append(s.models, SessionConfigOptionValue{Value: model, Name: model})
709+
}
574710
s.model = model
575711
s.mu.Unlock()
576712
}
@@ -581,6 +717,24 @@ func (s *acpSession) currentModel() string {
581717
return s.model
582718
}
583719

720+
func (s *acpSession) hasModel(model string) bool {
721+
model = strings.TrimSpace(model)
722+
s.mu.Lock()
723+
defer s.mu.Unlock()
724+
for _, option := range s.models {
725+
if option.Value == model {
726+
return true
727+
}
728+
}
729+
return false
730+
}
731+
732+
func (s *acpSession) configState() (string, []SessionConfigOptionValue, agent.PermissionMode) {
733+
s.mu.Lock()
734+
defer s.mu.Unlock()
735+
return s.model, append([]SessionConfigOptionValue(nil), s.models...), s.mode
736+
}
737+
584738
func (s *acpSession) appendHistory(rec turnRecord) {
585739
s.mu.Lock()
586740
s.history = append(s.history, rec)

0 commit comments

Comments
 (0)