@@ -7,6 +7,7 @@ import (
12 "github.com/charmbracelet/catwalk/pkg/catwalk"
13 "github.com/charmbracelet/catwalk/pkg/catwalk"
@@ -57,13 +58,12 @@ type Model struct {
59type sessionAgent struct {
60type sessionAgent struct {
60− largeModel Model
61+ largeModel Model
61− smallModel Model
62+ smallModel Model
62− systemPrompt string
63+ systemPrompt string
63− tools []ai.AgentTool
64+ tools []ai.AgentTool
64− maxOutputTokens int64
65− sessions session.Service
65+ sessions session.Service
66− messages message.Service
66+ messages message.Service
68 messageQueue *csync.Map[string, []SessionAgentCall]
68 messageQueue *csync.Map[string, []SessionAgentCall]
69 activeRequests *csync.Map[string, context.CancelFunc]
69 activeRequests *csync.Map[string, context.CancelFunc]
@@ -71,8 +71,24 @@ type sessionAgent struct {
72type SessionAgentOption func(*sessionAgent)
72type SessionAgentOption func(*sessionAgent)
74−func NewSessionAgent() SessionAgent {
74+func NewSessionAgent(
78+ sessions session.Service,
79+ messages message.Service,
80+ tools ...ai.AgentTool,
75− return &sessionAgent{}
82+ return &sessionAgent{
83+ largeModel: largeModel,
84+ smallModel: smallModel,
85+ systemPrompt: systemPrompt,
89+ messageQueue: csync.NewMap[string, []SessionAgentCall](),
90+ activeRequests: csync.NewMap[string, context.CancelFunc](),
78func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.AgentResult, error) {
94func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.AgentResult, error) {
@@ -103,7 +119,6 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
103 a.largeModel.model,
119 a.largeModel.model,
104 ai.WithSystemPrompt(a.systemPrompt),
120 ai.WithSystemPrompt(a.systemPrompt),
105 ai.WithTools(a.tools...),
121 ai.WithTools(a.tools...),
106− ai.WithMaxOutputTokens(a.maxOutputTokens),
109 currentSession, err := a.sessions.Get(ctx, call.SessionID)
124 currentSession, err := a.sessions.Get(ctx, call.SessionID)
@@ -116,9 +131,12 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
116 return nil, fmt.Errorf("failed to get session messages: %w", err)
131 return nil, fmt.Errorf("failed to get session messages: %w", err)
134+ var wg sync.WaitGroup
119 // Generate title if first message
135 // Generate title if first message
120 if len(msgs) == 0 {
136 if len(msgs) == 0 {
121− go a.generateTitle(ctx, currentSession, call.Prompt)
138+ a.generateTitle(ctx, currentSession, call.Prompt)
124 // Add the user message to the session
142 // Add the user message to the session
@@ -325,6 +343,7 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
325 if err != nil {
343 if err != nil {
326 return nil, err
344 return nil, err
329 queuedMessages, ok := a.messageQueue.Get(call.SessionID)
348 queuedMessages, ok := a.messageQueue.Get(call.SessionID)
330 if !ok || len(queuedMessages) == 0 {
349 if !ok || len(queuedMessages) == 0 {