diff --git a/agent/agent.go b/agent/agent.go index 73753c5..0a60de9 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -29,7 +29,7 @@ type Agent struct { agentsDir string baseLayers []string // system prompt layers for /agent reloads promptEnvExtra []string // PRIME_* vars for prompt resolution - newDispatcher func() tools.Dispatcher + newToolServer func() tools.Server newBackend func(string) (backend.Backend, error) currentAction atomic.Pointer[actionHandle] warnedContext bool @@ -258,10 +258,10 @@ func (ag *Agent) ToolCallCount() int64 { // SetSessionEnv injects session env vars into the execute server. func (ag *Agent) SetSessionEnv(sessionID string) { - if ag.runtime.Dispatcher == nil { + if ag.runtime.ExecServer == nil { return } - if srv, ok := ag.runtime.Dispatcher.GetServer("execute"); ok { + if srv := ag.runtime.ExecServer; srv != nil { if es, ok := srv.(tools.EnvSetter); ok { es.SetEnv("OLLIE_SESSION_ID", sessionID) if ag.id != "" { @@ -273,10 +273,10 @@ func (ag *Agent) SetSessionEnv(sessionID string) { // SetEnv stores an environment variable on the execute server. func (ag *Agent) SetEnv(key, value string) { - if ag.runtime.Dispatcher == nil { + if ag.runtime.ExecServer == nil { return } - if srv, ok := ag.runtime.Dispatcher.GetServer("execute"); ok { + if srv := ag.runtime.ExecServer; srv != nil { if es, ok := srv.(tools.EnvSetter); ok { es.SetEnv(key, value) } @@ -285,10 +285,10 @@ func (ag *Agent) SetEnv(key, value string) { // Close releases agent resources (dispatcher, execute server). func (ag *Agent) Close() { - if ag.runtime.Dispatcher == nil { + if ag.runtime.ExecServer == nil { return } - if srv, ok := ag.runtime.Dispatcher.GetServer("execute"); ok { + if srv := ag.runtime.ExecServer; srv != nil { if c, ok := srv.(interface{ Close() }); ok { c.Close() } @@ -302,8 +302,8 @@ func (ag *Agent) SetCWD(dir string) { if oldCwd != "" && dir != "" && oldCwd != dir { ag.runtime.Preamble = strings.ReplaceAll(ag.runtime.Preamble, oldCwd, dir) } - if ag.runtime != nil && ag.runtime.Dispatcher != nil { - if srv, ok := ag.runtime.Dispatcher.GetServer("execute"); ok { + if ag.runtime.ExecServer != nil { + if srv := ag.runtime.ExecServer; srv != nil { if ws, ok := srv.(tools.CWDSetter); ok { ws.SetCWD(dir) } @@ -472,13 +472,9 @@ func (ag *Agent) React(responseID, emoji string) error { return nil } -// execServer returns the execute server interface, or nil if unavailable. -func (ag *Agent) execServer() interface{} { - if ag.runtime.Dispatcher == nil { - return nil - } - srv, _ := ag.runtime.Dispatcher.GetServer("execute") - return srv +// execServer returns the execute server, or nil if unavailable. +func (ag *Agent) execServer() tools.Server { + return ag.runtime.ExecServer } // Queue pushes a prompt onto the agent's FIFO. diff --git a/agent/build_runtime.go b/agent/build_runtime.go index 35024b9..bfb53d9 100644 --- a/agent/build_runtime.go +++ b/agent/build_runtime.go @@ -17,7 +17,7 @@ import ( // env provides additional environment variables injected into prompt resolution // subprocesses (e.g. OLLIE_SESSION_ID=xxx). // The caller is responsible for registering all servers on d before calling this. -func BuildRuntime(cfg *AgentConfig, d tools.Dispatcher, cwd string, env []string, baseLayers ...string) *Runtime { +func BuildRuntime(cfg *AgentConfig, srv tools.Server, cwd string, env []string, baseLayers ...string) *Runtime { var messages []string var allToolInfos []tools.ToolInfo @@ -25,7 +25,7 @@ func BuildRuntime(cfg *AgentConfig, d tools.Dispatcher, cwd string, env []string if cfg == nil || cfg.ToolsEnabled() { var listErr error - allToolInfos, listErr = d.ListTools() + allToolInfos, listErr = srv.ListTools() if listErr != nil { messages = append(messages, fmt.Sprintf("list tools: %v", listErr)) } @@ -69,30 +69,14 @@ func BuildRuntime(cfg *AgentConfig, d tools.Dispatcher, cwd string, env []string } maxSteps = cfg.MaxSteps if len(cfg.AllowTools) > 0 { - if srv, ok := d.GetServer("execute"); ok { - if rs, ok := srv.(tools.ToolRestrictionSetter); ok { - rs.SetAllowTools(cfg.AllowTools) - } + if rs, ok := srv.(tools.ToolRestrictionSetter); ok { + rs.SetAllowTools(cfg.AllowTools) } } } exec := func(ctx context.Context, name string, args json.RawMessage) (string, []backend.ContentBlock, error) { - infos, listErr := d.ListTools() - if listErr != nil { - return "", nil, listErr - } - server := "" - for _, t := range infos { - if t.Name == name { - server = t.Server - break - } - } - if server == "" { - return "", nil, fmt.Errorf("unknown tool: %s", name) - } - raw, err := d.Dispatch(ctx, server, name, args) + raw, err := srv.CallTool(ctx, name, args) if err != nil { return "", nil, err } @@ -104,24 +88,20 @@ func BuildRuntime(cfg *AgentConfig, d tools.Dispatcher, cwd string, env []string } var classify toolClassifier - if srv, ok := d.GetServer("execute"); ok { - if pc, ok := srv.(tools.ParallelClassifier); ok { - classify = pc.IsParallelRead - } + if pc, ok := srv.(tools.ParallelClassifier); ok { + classify = pc.IsParallelRead } var tierFn func(string, json.RawMessage) ResultTier - if srv, ok := d.GetServer("execute"); ok { - if tc, ok := srv.(tools.TierClassifier); ok { - tierFn = func(name string, args json.RawMessage) ResultTier { - switch tc.ResultTierArgs(name, args) { - case "cold": - return TierCold - case "warm": - return TierWarm - default: - return TierHot - } + if tc, ok := srv.(tools.TierClassifier); ok { + tierFn = func(name string, args json.RawMessage) ResultTier { + switch tc.ResultTierArgs(name, args) { + case "cold": + return TierCold + case "warm": + return TierWarm + default: + return TierHot } } } @@ -160,7 +140,7 @@ func BuildRuntime(cfg *AgentConfig, d tools.Dispatcher, cwd string, env []string } return &Runtime{ - Dispatcher: d, + ExecServer: srv, Tools: allTools, Exec: exec, ClassifyTool: classify, diff --git a/agent/commands.go b/agent/commands.go index 5922caf..e0af3db 100644 --- a/agent/commands.go +++ b/agent/commands.go @@ -180,7 +180,7 @@ func (ag *Agent) cmdAgent(_ context.Context, args []string) { ag.emit(infoEvent(fmt.Sprintf("error: agent %q: %v", name, err))) return } - disp := ag.newDispatcher() + disp := ag.newToolServer() env := []string{"OLLIE_SESSION_ID=" + ag.sessionID, "OLLIE_UNAME=" + ag.id} env = append(env, ag.promptEnvExtra...) rt := BuildRuntime(cfg, disp, ag.cwd, env, ag.baseLayers...) diff --git a/agent/new.go b/agent/new.go index 9533447..a94b434 100644 --- a/agent/new.go +++ b/agent/new.go @@ -19,7 +19,7 @@ type AgentCfg struct { CWD string // working directory for tool execution BaseLayers []string PromptEnvExtra []string - NewDispatcher func() tools.Dispatcher + NewToolServer func() tools.Server NewBackend func(string) (backend.Backend, error) Bus *pubsub.Bus Log *olog.Logger @@ -42,7 +42,7 @@ func NewAgent(cfg AgentCfg) *Agent { cwd: cfg.CWD, baseLayers: cfg.BaseLayers, promptEnvExtra: cfg.PromptEnvExtra, - newDispatcher: cfg.NewDispatcher, + newToolServer: cfg.NewToolServer, newBackend: cfg.NewBackend, bus: cfg.Bus, log: cfg.Log, diff --git a/agent/runtime.go b/agent/runtime.go index 5074a05..ba3f0ee 100644 --- a/agent/runtime.go +++ b/agent/runtime.go @@ -13,7 +13,7 @@ import ( // agents replaces it atomically. type Runtime struct { Backend backend.Backend - Dispatcher tools.Dispatcher + ExecServer tools.Server // the execute server (tool runtime, env, cwd) Hooks Hooks Preamble string // compiled system prompt Tools []backend.Tool diff --git a/session/session.go b/session/session.go index 0066411..72f2ceb 100644 --- a/session/session.go +++ b/session/session.go @@ -32,7 +32,7 @@ type Config struct { CWD string History *agent.History Runtime *agent.Runtime - NewDispatcher func() tools.Dispatcher + NewToolServer func() tools.Server NewBackend func(string) (backend.Backend, error) Log *olog.Logger MaxSteps int @@ -150,7 +150,7 @@ func New(cfg Config) *Session { CWD: paths.ExpandHome(cfg.CWD), BaseLayers: cfg.BaseLayers, PromptEnvExtra: cfg.PromptEnvExtra, - NewDispatcher: cfg.NewDispatcher, + NewToolServer: cfg.NewToolServer, NewBackend: cfg.NewBackend, Bus: bus, Log: log,