agent.go

  1package backend
  2
  3import (
  4	"context"
  5
  6	"github.com/charmbracelet/crush/internal/config"
  7	"github.com/charmbracelet/crush/internal/proto"
  8)
  9
 10// SendMessage sends a prompt to the agent coordinator for the given
 11// workspace and session.
 12func (b *Backend) SendMessage(ctx context.Context, workspaceID string, msg proto.AgentMessage) error {
 13	ws, err := b.GetWorkspace(workspaceID)
 14	if err != nil {
 15		return err
 16	}
 17
 18	if ws.AgentCoordinator == nil {
 19		return ErrAgentNotInitialized
 20	}
 21
 22	_, err = ws.AgentCoordinator.Run(ctx, msg.SessionID, msg.Prompt)
 23	return err
 24}
 25
 26// GetAgentInfo returns the agent's model and busy status.
 27func (b *Backend) GetAgentInfo(workspaceID string) (proto.AgentInfo, error) {
 28	ws, err := b.GetWorkspace(workspaceID)
 29	if err != nil {
 30		return proto.AgentInfo{}, err
 31	}
 32
 33	var agentInfo proto.AgentInfo
 34	if ws.AgentCoordinator != nil {
 35		m := ws.AgentCoordinator.Model()
 36		agentInfo = proto.AgentInfo{
 37			Model:    m.CatwalkCfg,
 38			ModelCfg: m.ModelCfg,
 39			IsBusy:   ws.AgentCoordinator.IsBusy(),
 40			IsReady:  true,
 41		}
 42	}
 43	return agentInfo, nil
 44}
 45
 46// InitAgent initializes the coder agent for the workspace.
 47func (b *Backend) InitAgent(ctx context.Context, workspaceID string) error {
 48	ws, err := b.GetWorkspace(workspaceID)
 49	if err != nil {
 50		return err
 51	}
 52
 53	return ws.InitCoderAgent(ctx)
 54}
 55
 56// UpdateAgent reloads the agent model configuration.
 57func (b *Backend) UpdateAgent(ctx context.Context, workspaceID string) error {
 58	ws, err := b.GetWorkspace(workspaceID)
 59	if err != nil {
 60		return err
 61	}
 62
 63	return ws.UpdateAgentModel(ctx)
 64}
 65
 66// CancelSession cancels an ongoing agent operation for the given
 67// session.
 68func (b *Backend) CancelSession(workspaceID, sessionID string) error {
 69	ws, err := b.GetWorkspace(workspaceID)
 70	if err != nil {
 71		return err
 72	}
 73
 74	if ws.AgentCoordinator != nil {
 75		ws.AgentCoordinator.Cancel(sessionID)
 76	}
 77	return nil
 78}
 79
 80// SummarizeSession triggers a session summarization.
 81func (b *Backend) SummarizeSession(ctx context.Context, workspaceID, sessionID string) error {
 82	ws, err := b.GetWorkspace(workspaceID)
 83	if err != nil {
 84		return err
 85	}
 86
 87	if ws.AgentCoordinator == nil {
 88		return ErrAgentNotInitialized
 89	}
 90
 91	return ws.AgentCoordinator.Summarize(ctx, sessionID)
 92}
 93
 94// QueuedPrompts returns the number of queued prompts for the session.
 95func (b *Backend) QueuedPrompts(workspaceID, sessionID string) (int, error) {
 96	ws, err := b.GetWorkspace(workspaceID)
 97	if err != nil {
 98		return 0, err
 99	}
100
101	if ws.AgentCoordinator == nil {
102		return 0, nil
103	}
104
105	return ws.AgentCoordinator.QueuedPrompts(sessionID), nil
106}
107
108// ClearQueue clears the prompt queue for the session.
109func (b *Backend) ClearQueue(workspaceID, sessionID string) error {
110	ws, err := b.GetWorkspace(workspaceID)
111	if err != nil {
112		return err
113	}
114
115	if ws.AgentCoordinator != nil {
116		ws.AgentCoordinator.ClearQueue(sessionID)
117	}
118	return nil
119}
120
121// QueuedPromptsList returns the list of queued prompt strings for a
122// session.
123func (b *Backend) QueuedPromptsList(workspaceID, sessionID string) ([]string, error) {
124	ws, err := b.GetWorkspace(workspaceID)
125	if err != nil {
126		return nil, err
127	}
128
129	if ws.AgentCoordinator == nil {
130		return nil, nil
131	}
132
133	return ws.AgentCoordinator.QueuedPromptsList(sessionID), nil
134}
135
136// GetDefaultSmallModel returns the default small model for a provider.
137func (b *Backend) GetDefaultSmallModel(workspaceID, providerID string) (config.SelectedModel, error) {
138	ws, err := b.GetWorkspace(workspaceID)
139	if err != nil {
140		return config.SelectedModel{}, err
141	}
142
143	return ws.GetDefaultSmallModel(providerID), nil
144}