session: wire registries via Options instead of post-hoc

Pass WithToolRegistry/WithSkillsRegistry into execute.Decl() so every
Server the factory creates has registries set. Removes fragile post-
construction wiring that only applied to the first dispatcher instance.
This commit is contained in:
Levi Neely 2026-07-29 10:50:26 +02:00
parent 0a9392e5fd
commit f0942e81ba
1 changed files with 12 additions and 22 deletions

View File

@ -872,6 +872,12 @@ func (s *Manager) restoreSession(ps *agent.PersistedSession) error {
if s.cfg.Yolo {
execOpts = append(execOpts, execute.WithYolo())
}
if s.cfg.ToolRegistry != nil {
execOpts = append(execOpts, execute.WithToolRegistry(s.cfg.ToolRegistry, sessID))
}
if s.cfg.SkillsRegistry != nil {
execOpts = append(execOpts, execute.WithSkillsRegistry(s.cfg.SkillsRegistry))
}
newDisp = tools.NewDispatcherFunc(map[string]func() tools.Server{
"execute": execute.Decl(cwd, execOpts...),
})
@ -909,17 +915,6 @@ func (s *Manager) restoreSession(ps *agent.PersistedSession) error {
envBlock := prompts.Environment(cwd, platform, isGitRepo, "")
disp := newDisp()
// Register the registry-backed tool server for lazy tool promotion
if s.cfg.ToolRegistry != nil {
if execSrv, ok := disp.GetServer("execute"); ok {
if es, ok := execSrv.(*execute.Server); ok {
es.SetToolRegistry(s.cfg.ToolRegistry, sessID)
if s.cfg.SkillsRegistry != nil {
es.SetSkillsRegistry(s.cfg.SkillsRegistry)
}
}
}
}
rt := agent.BuildRuntime(cfg, disp, cwd, env, sysPrompt, opModel, envBlock)
restoredSession := agent.RestoreSession(ps)
@ -1168,6 +1163,12 @@ func (s *Manager) CreateSession(args []string) (string, error) {
if cfg != nil && len(cfg.AllowTools) > 0 {
execOpts = append(execOpts, execute.WithAllowTools(cfg.AllowTools))
}
if s.cfg.ToolRegistry != nil {
execOpts = append(execOpts, execute.WithToolRegistry(s.cfg.ToolRegistry, sessID))
}
if s.cfg.SkillsRegistry != nil {
execOpts = append(execOpts, execute.WithSkillsRegistry(s.cfg.SkillsRegistry))
}
var newDisp func() tools.Dispatcher
var remoteEnv []string // extra env vars from remote host info
if remoteTarget != "" {
@ -1236,17 +1237,6 @@ func (s *Manager) CreateSession(args []string) (string, error) {
envBlock := prompts.Environment(cwd, platform, isGitRepo, "")
disp := newDisp()
// Register the registry-backed tool server for lazy tool promotion
if s.cfg.ToolRegistry != nil {
if execSrv, ok := disp.GetServer("execute"); ok {
if es, ok := execSrv.(*execute.Server); ok {
es.SetToolRegistry(s.cfg.ToolRegistry, sessID)
if s.cfg.SkillsRegistry != nil {
es.SetSkillsRegistry(s.cfg.SkillsRegistry)
}
}
}
}
rt := agent.BuildRuntime(cfg, disp, cwd, env, sysPrompt, opModel, envBlock)
// sessPtr is set after NewSession; the ReadPlanStep closure captures it.