From a2de021f3c6ae54ddb5f6f0b343f8526a01c1efa Mon Sep 17 00:00:00 2001 From: Ian Gulliver Date: Tue, 20 Oct 2026 22:53:47 -0700 Subject: [PATCH] Unify model changes and shell command argument handling --- internal/agents/bash.go | 26 +++++++++++++++++--------- internal/agents/clients.go | 17 +++++------------ internal/agents/lead.go | 13 +++---------- internal/agents/manager.go | 13 +++---------- internal/agents/process.go | 16 ++++++++++++---- internal/agents/queue.go | 13 +++---------- 6 files changed, 43 insertions(+), 55 deletions(-) diff --git a/internal/agents/bash.go b/internal/agents/bash.go index 49648c1..81ba709 100644 --- a/internal/agents/bash.go +++ b/internal/agents/bash.go @@ -21,6 +21,20 @@ type bashArgs struct { TimeoutMs *int `json:"timeout_ms"` } +func (a bashArgs) timeout() (time.Duration, error) { + if strings.TrimSpace(a.Command) == "" { + return 0, errors.New("command is empty") + } + timeout := defaultBashTimeout + if a.TimeoutMs != nil { + timeout = time.Duration(*a.TimeoutMs) * time.Millisecond + } + if timeout <= 0 || timeout > maxBashTimeout { + return 0, fmt.Errorf("timeout_ms must be between 1 and %d", maxBashTimeout.Milliseconds()) + } + return timeout, nil +} + func (m *Manager) bashTool(l *live) mcp.Tool { return mcp.Tool{ Name: "bash", @@ -47,15 +61,9 @@ func (m *Manager) bash(ctx context.Context, l *live, raw json.RawMessage) (strin if err != nil { return "", fmt.Errorf("malformed bash call: %w", err) } - if strings.TrimSpace(args.Command) == "" { - return "", errors.New("command is empty") - } - timeout := defaultBashTimeout - if args.TimeoutMs != nil { - timeout = time.Duration(*args.TimeoutMs) * time.Millisecond - } - if timeout <= 0 || timeout > maxBashTimeout { - return "", fmt.Errorf("timeout_ms must be between 1 and %d", maxBashTimeout.Milliseconds()) + timeout, err := args.timeout() + if err != nil { + return "", err } task, err := m.store.GetTask(ctx, l.taskID) if err != nil { diff --git a/internal/agents/clients.go b/internal/agents/clients.go index 4505c64..025db4b 100644 --- a/internal/agents/clients.go +++ b/internal/agents/clients.go @@ -256,9 +256,8 @@ func (m *Manager) listClients() string { } type clientBashArgs struct { - Client string `json:"client"` - Command string `json:"command"` - TimeoutMs *int `json:"timeout_ms"` + Client string `json:"client"` + bashArgs } func (m *Manager) clientBash(ctx context.Context, raw json.RawMessage) (string, error) { @@ -267,15 +266,9 @@ func (m *Manager) clientBash(ctx context.Context, raw json.RawMessage) (string, if err != nil { return "", fmt.Errorf("malformed client_bash call: %w", err) } - if strings.TrimSpace(args.Command) == "" { - return "", errors.New("command is empty") - } - timeout := defaultBashTimeout - if args.TimeoutMs != nil { - timeout = time.Duration(*args.TimeoutMs) * time.Millisecond - } - if timeout <= 0 || timeout > maxBashTimeout { - return "", fmt.Errorf("timeout_ms must be between 1 and %d", maxBashTimeout.Milliseconds()) + timeout, err := args.timeout() + if err != nil { + return "", err } exitCode, output, err := m.runClientBash(ctx, args.Client, args.Command, timeout) diff --git a/internal/agents/lead.go b/internal/agents/lead.go index c60b1bb..afa5199 100644 --- a/internal/agents/lead.go +++ b/internal/agents/lead.go @@ -147,16 +147,9 @@ func (m *Manager) MessageLead(ctx context.Context, user api.User, text, model, s } } - if model != "" && model != sess.Model { - err = checkModel(sess, model) - if err != nil { - return api.AgentSession{}, err - } - err = m.store.SetSessionModel(ctx, sess.Id, model, user.Login) - if err != nil { - return api.AgentSession{}, err - } - sess.Model = model + sess, err = m.changeModel(ctx, sess, model, user.Login) + if err != nil { + return api.AgentSession{}, err } err = m.store.AddLeadInput(ctx, store.LeadInput{Owner: user.Login, Kind: store.LeadMessage, Actor: user.Login, Text: text, Images: images, SendID: sendID}) if err != nil { diff --git a/internal/agents/manager.go b/internal/agents/manager.go index e21462e..d94ed75 100644 --- a/internal/agents/manager.go +++ b/internal/agents/manager.go @@ -262,16 +262,9 @@ func (m *Manager) Send(ctx context.Context, sessionID, actor, text, model, sendI return api.AgentSession{}, fmt.Errorf("%w: session %s is %s", ErrConflict, sess.Id, sess.State) } - if model != "" && model != sess.Model { - err = checkModel(sess, model) - if err != nil { - return api.AgentSession{}, err - } - err = m.store.SetSessionModel(ctx, sess.Id, model, actor) - if err != nil { - return api.AgentSession{}, err - } - sess.Model = model + sess, err = m.changeModel(ctx, sess, model, actor) + if err != nil { + return api.AgentSession{}, err } task, err := m.taskFor(ctx, sess) diff --git a/internal/agents/process.go b/internal/agents/process.go index e82baed..b35073f 100644 --- a/internal/agents/process.go +++ b/internal/agents/process.go @@ -50,11 +50,19 @@ func initialVendorSessionID(vendor api.AgentSessionVendor, id string) string { return id } -func checkModel(sess api.AgentSession, model string) error { - if vendorOf(model) != sess.Vendor { - return fmt.Errorf("%w: session %s is %s; %s is a %s model", ErrInvalid, sess.Id, sess.Vendor, model, vendorOf(model)) +func (m *Manager) changeModel(ctx context.Context, sess api.AgentSession, model, actor string) (api.AgentSession, error) { + if model == "" || model == sess.Model { + return sess, nil } - return nil + if vendorOf(model) != sess.Vendor { + return api.AgentSession{}, fmt.Errorf("%w: session %s is %s; %s is a %s model", ErrInvalid, sess.Id, sess.Vendor, model, vendorOf(model)) + } + err := m.store.SetSessionModel(ctx, sess.Id, model, actor) + if err != nil { + return api.AgentSession{}, err + } + sess.Model = model + return sess, nil } type claudeProcess struct { diff --git a/internal/agents/queue.go b/internal/agents/queue.go index 904dadb..f1a1165 100644 --- a/internal/agents/queue.go +++ b/internal/agents/queue.go @@ -64,16 +64,9 @@ func (m *Manager) Message(ctx context.Context, sessionID, actor, text, model, se return sess, false, err } - if model != "" && model != sess.Model { - err = checkModel(sess, model) - if err != nil { - return api.AgentSession{}, false, err - } - err = m.store.SetSessionModel(ctx, sess.Id, model, actor) - if err != nil { - return api.AgentSession{}, false, err - } - sess.Model = model + sess, err = m.changeModel(ctx, sess, model, actor) + if err != nil { + return api.AgentSession{}, false, err } task, err := m.store.GetTask(ctx, sess.TaskId)