From f0942e81ba2fc15cecd8eab9404853074041eb06 Mon Sep 17 00:00:00 2001 From: Levi Neely Date: Wed, 29 Jul 2026 10:50:26 +0200 Subject: [PATCH] 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. --- session/session.go | 34 ++++++++++++---------------------- 1 file changed, 12 insertions(+), 22 deletions(-) diff --git a/session/session.go b/session/session.go index bfe2e05..7e28d56 100644 --- a/session/session.go +++ b/session/session.go @@ -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.