@@ -4,8 +4,6 @@ import (
@@ -16,133 +14,101 @@ import (
16 "github.com/kujtimiihoxha/termai/internal/llm/tools"
14 "github.com/kujtimiihoxha/termai/internal/llm/tools"
17 "github.com/kujtimiihoxha/termai/internal/logging"
15 "github.com/kujtimiihoxha/termai/internal/logging"
18 "github.com/kujtimiihoxha/termai/internal/message"
16 "github.com/kujtimiihoxha/termai/internal/message"
17+ "github.com/kujtimiihoxha/termai/internal/permission"
19 "github.com/kujtimiihoxha/termai/internal/session"
18 "github.com/kujtimiihoxha/termai/internal/session"
22// Common errors
21// Common errors
24− ErrProviderNotEnabled = errors.New("provider is not enabled")
25− ErrRequestCancelled = errors.New("request cancelled by user")
23+ ErrRequestCancelled = errors.New("request cancelled by user")
26− ErrSessionBusy = errors.New("session is currently processing another request")
24+ ErrSessionBusy = errors.New("session is currently processing another request")
29−// Service defines the interface for generating responses
27+type AgentEvent struct {
28+ message message.Message
32+func (e *AgentEvent) Err() error {
36+func (e *AgentEvent) Response() message.Message {
30type Service interface {
40type Service interface {
31− Generate(ctx context.Context, sessionID string, content string) error
41+ Run(ctx context.Context, sessionID string, content string) (<-chan AgentEvent, error)
32− Cancel(sessionID string) error
42+ Cancel(sessionID string)
43+ IsSessionBusy(sessionID string) bool
35type agent struct {
46type agent struct {
36− sessions session.Service
47+ sessions session.Service
37− messages message.Service
48+ messages message.Service
38− model models.Model
49+
39− tools []tools.BaseTool
50+ tools []tools.BaseTool
40− agent provider.Provider
51+ provider provider.Provider
41− titleGenerator provider.Provider
53+ titleProvider provider.Provider
42− activeRequests sync.Map // map[sessionID]context.CancelFunc
55+ activeRequests sync.Map
45−// NewAgent creates a new agent instance with the given model and tools
58+func NewAgent(
46−func NewAgent(ctx context.Context, sessions session.Service, messages message.Service, model models.Model, tools []tools.BaseTool) (Service, error) {
59+ agentName config.AgentName,
47− agentProvider, titleGenerator, err := getAgentProviders(ctx, model)
60+ sessions session.Service,
61+ messages message.Service,
62+ agentTools []tools.BaseTool,
64+ agentProvider, err := createAgentProvider(agentName)
48 if err != nil {
65 if err != nil {
49− return nil, fmt.Errorf("failed to initialize providers: %w", err)
66+ return nil, err
68+ var titleProvider provider.Provider
69+ // Only generate titles for the coder agent
70+ if agentName == config.AgentCoder {
71+ titleProvider, err = createAgentProvider(config.AgentTitle)
52− return &agent{
77+ agent := &agent{
53− model: model,
78+ provider: agentProvider,
56 messages: messages,
79 messages: messages,
57− agent: agentProvider,
80+ sessions: sessions,
58− titleGenerator: titleGenerator,
81+ tools: agentTools,
82+ titleProvider: titleProvider,
59 activeRequests: sync.Map{},
83 activeRequests: sync.Map{},
63−// Cancel cancels an active request by session ID
64−func (a *agent) Cancel(sessionID string) error {
89+func (a *agent) Cancel(sessionID string) {
65 if cancelFunc, exists := a.activeRequests.LoadAndDelete(sessionID); exists {
90 if cancelFunc, exists := a.activeRequests.LoadAndDelete(sessionID); exists {
66 if cancel, ok := cancelFunc.(context.CancelFunc); ok {
91 if cancel, ok := cancelFunc.(context.CancelFunc); ok {
67 logging.InfoPersist(fmt.Sprintf("Request cancellation initiated for session: %s", sessionID))
92 logging.InfoPersist(fmt.Sprintf("Request cancellation initiated for session: %s", sessionID))
72− return errors.New("no active request found for this session")
75−// Generate starts the generation process
76−func (a *agent) Generate(ctx context.Context, sessionID string, content string) error {
77− // Check if this session already has an active request
78− if _, busy := a.activeRequests.Load(sessionID); busy {
79− return ErrSessionBusy
82− // Create a cancellable context
83− genCtx, cancel := context.WithCancel(ctx)
85− // Store cancel function to allow user cancellation
86− a.activeRequests.Store(sessionID, cancel)
88− // Launch the generation in a goroutine
91− if r := recover(); r != nil {
92− logging.ErrorPersist(fmt.Sprintf("Panic in Generate: %v", r))
94− // dump stack trace into a file
95− file, err := os.Create("panic.log")
97− logging.ErrorPersist(fmt.Sprintf("Failed to create panic log: %v", err))
103− stackTrace := debug.Stack()
104− if _, err := file.Write(stackTrace); err != nil {
105− logging.ErrorPersist(fmt.Sprintf("Failed to write panic log: %v", err))
110− defer a.activeRequests.Delete(sessionID)
113− if err := a.generate(genCtx, sessionID, content); err != nil {
114− if !errors.Is(err, ErrRequestCancelled) && !errors.Is(err, context.Canceled) {
115− // Log the error (avoid logging cancellations as they're expected)
116− logging.ErrorPersist(fmt.Sprintf("Generation error for session %s: %v", sessionID, err))
118− // You may want to create an error message in the chat
119− bgCtx := context.Background()
120− errorMsg := fmt.Sprintf("Sorry, an error occurred: %v", err)
121− _, createErr := a.messages.Create(bgCtx, sessionID, message.CreateMessageParams{
122− Role: message.System,
123− Parts: []message.ContentPart{
124− message.TextContent{
129− if createErr != nil {
130− logging.ErrorPersist(fmt.Sprintf("Failed to create error message: %v", createErr))
139−// IsSessionBusy checks if a session currently has an active request
140func (a *agent) IsSessionBusy(sessionID string) bool {
98func (a *agent) IsSessionBusy(sessionID string) bool {
141 _, busy := a.activeRequests.Load(sessionID)
99 _, busy := a.activeRequests.Load(sessionID)
142 return busy
100 return busy
143−} // handleTitleGeneration asynchronously generates a title for new sessions
101+}
144−func (a *agent) handleTitleGeneration(ctx context.Context, sessionID, content string) {
103+func (a *agent) generateTitle(ctx context.Context, sessionID string, content string) error {
104+ if a.titleProvider == nil {
107+ session, err := a.sessions.Get(ctx, sessionID)
145− response, err := a.titleGenerator.SendMessages(
111+ response, err := a.titleProvider.SendMessages(
147 []message.Message{
113 []message.Message{
@@ -154,121 +120,152 @@ func (a *agent) handleTitleGeneration(ctx context.Context, sessionID, content st
157− nil,
123+ make([]tools.BaseTool, 0),
159 if err != nil {
125 if err != nil {
160− logging.ErrorPersist(fmt.Sprintf("Failed to generate title: %v", err))
126+ return err
164− session, err := a.sessions.Get(ctx, sessionID)
129+ title := strings.TrimSpace(strings.ReplaceAll(response.Content, "\n", " "))
165− if err != nil {
130+ if title == "" {
166− logging.ErrorPersist(fmt.Sprintf("Failed to get session: %v", err))
167− return
131+ return nil
170− if response.Content != "" {
171− session.Title = strings.TrimSpace(response.Content)
134+ session.Title = title
172− session.Title = strings.ReplaceAll(session.Title, "\n", " ")
173− if _, err := a.sessions.Save(ctx, session); err != nil {
135+ _, err = a.sessions.Save(ctx, session)
174− logging.ErrorPersist(fmt.Sprintf("Failed to save session title: %v", err))
136+ return err
139+func (a *agent) err(err error) AgentEvent {
179−// TrackUsage updates token usage statistics for the session
180−func (a *agent) TrackUsage(ctx context.Context, sessionID string, model models.Model, usage provider.TokenUsage) error {
145+func (a *agent) Run(ctx context.Context, sessionID string, content string) (<-chan AgentEvent, error) {
181− session, err := a.sessions.Get(ctx, sessionID)
146+ events := make(chan AgentEvent)
182− if err != nil {
147+ if a.IsSessionBusy(sessionID) {
183− return fmt.Errorf("failed to get session: %w", err)
148+ return nil, ErrSessionBusy
186− cost := model.CostPer1MInCached/1e6*float64(usage.CacheCreationTokens) +
151+ genCtx, cancel := context.WithCancel(ctx)
187− model.CostPer1MOutCached/1e6*float64(usage.CacheReadTokens) +
152+
188− model.CostPer1MIn/1e6*float64(usage.InputTokens) +
153+ a.activeRequests.Store(sessionID, cancel)
189− model.CostPer1MOut/1e6*float64(usage.OutputTokens)
154+ go func() {
155+ logging.Debug("Request started", "sessionID", sessionID)
156+ defer logging.RecoverPanic("agent.Run", func() {
157+ events <- a.err(fmt.Errorf("panic while running the agent"))
191− session.Cost += cost
160+ result := a.processGeneration(genCtx, sessionID, content)
192− session.CompletionTokens += usage.OutputTokens
161+ if result.Err() != nil && !errors.Is(result.Err(), ErrRequestCancelled) && !errors.Is(result.Err(), context.Canceled) {
193− session.PromptTokens += usage.InputTokens
162+ logging.ErrorPersist(fmt.Sprintf("Generation error for session %s: %v", sessionID, result))
164+ logging.Debug("Request completed", "sessionID", sessionID)
165+ a.activeRequests.Delete(sessionID)
195− _, err = a.sessions.Save(ctx, session)
173+func (a *agent) processGeneration(ctx context.Context, sessionID, content string) AgentEvent {
174+ // List existing messages; if none, start title generation asynchronously.
175+ msgs, err := a.messages.List(ctx, sessionID)
196 if err != nil {
176 if err != nil {
197− return fmt.Errorf("failed to save session: %w", err)
177+ return a.err(fmt.Errorf("failed to list messages: %w", err))
181+ defer logging.RecoverPanic("agent.Run", func() {
182+ logging.ErrorPersist("panic while generating title")
184+ titleErr := a.generateTitle(context.Background(), sessionID, content)
185+ if titleErr != nil {
186+ logging.ErrorPersist(fmt.Sprintf("failed to generate title: %v", titleErr))
202−// processEvent handles different types of events during generation
191+ userMsg, err := a.createUserMessage(ctx, sessionID, content)
203−func (a *agent) processEvent(
192+ if err != nil {
204− ctx context.Context,
193+ return a.err(fmt.Errorf("failed to create user message: %w", err))
206− assistantMsg *message.Message,
207− event provider.ProviderEvent,
213− // Continue processing
216− switch event.Type {
196+ // Append the new user message to the conversation history.
217− case provider.EventThinkingDelta:
197+ msgHistory := append(msgs, userMsg)
218− assistantMsg.AppendReasoningContent(event.Content)
198+ for {
219− return a.messages.Update(ctx, *assistantMsg)
199+ // Check for cancellation before each iteration
220− case provider.EventContentDelta:
200+ select {
221− assistantMsg.AppendContent(event.Content)
201+ case <-ctx.Done():
222− return a.messages.Update(ctx, *assistantMsg)
202+ return a.err(ctx.Err())
223− case provider.EventError:
203+ default:
224− if errors.Is(event.Error, context.Canceled) {
204+ // Continue processing
225− logging.InfoPersist(fmt.Sprintf("Event processing canceled for session: %s", sessionID))
226− return context.Canceled
228− logging.ErrorPersist(event.Error.Error())
206+ agentMessage, toolResults, err := a.streamAndHandleEvents(ctx, sessionID, msgHistory)
229− return event.Error
207+ if err != nil {
230− case provider.EventWarning:
208+ if errors.Is(err, context.Canceled) {
231− logging.WarnPersist(event.Info)
209+ return a.err(ErrRequestCancelled)
232− case provider.EventInfo:
210+ }
233− logging.InfoPersist(event.Info)
234− case provider.EventComplete:
235− assistantMsg.SetToolCalls(event.Response.ToolCalls)
236− assistantMsg.AddFinish(event.Response.FinishReason)
237− if err := a.messages.Update(ctx, *assistantMsg); err != nil {
238− return fmt.Errorf("failed to update message: %w", err)
211+ return a.err(fmt.Errorf("failed to process events: %w", err))
213+ logging.Info("Result", "message", agentMessage.FinishReason(), "toolResults", toolResults)
214+ if (agentMessage.FinishReason() == message.FinishReasonToolUse) && toolResults != nil {
215+ // We are not done, we need to respond with the tool response
216+ msgHistory = append(msgHistory, agentMessage, *toolResults)
220+ message: agentMessage,
240− return a.TrackUsage(ctx, sessionID, a.model, event.Response.Usage)
243− return nil
225+func (a *agent) createUserMessage(ctx context.Context, sessionID, content string) (message.Message, error) {
226+ return a.messages.Create(ctx, sessionID, message.CreateMessageParams{
228+ Parts: []message.ContentPart{
229+ message.TextContent{Text: content},
246−// ExecuteTools runs all tool calls sequentially and returns the results
234+func (a *agent) streamAndHandleEvents(ctx context.Context, sessionID string, msgHistory []message.Message) (message.Message, *message.Message, error) {
247−func (a *agent) ExecuteTools(ctx context.Context, toolCalls []message.ToolCall, tls []tools.BaseTool) ([]message.ToolResult, error) {
235+ eventChan := a.provider.StreamResponse(ctx, msgHistory, a.tools)
248− toolResults := make([]message.ToolResult, len(toolCalls))
236+
237+ assistantMsg, err := a.messages.Create(ctx, sessionID, message.CreateMessageParams{
238+ Role: message.Assistant,
239+ Parts: []message.ContentPart{},
240+ Model: a.provider.Model().ID,
243+ return assistantMsg, nil, fmt.Errorf("failed to create assistant message: %w", err)
250− // Create a child context that can be canceled
246+ // Add the session and message ID into the context if needed by tools.
251− ctx, cancel := context.WithCancel(ctx)
247+ ctx = context.WithValue(ctx, tools.MessageIDContextKey, assistantMsg.ID)
252− defer cancel()
248+ ctx = context.WithValue(ctx, tools.SessionIDContextKey, sessionID)
254− // Check if already canceled before starting any execution
250+ // Process each event in the stream.
255− if ctx.Err() != nil {
251+ for event := range eventChan {
256− // Mark all tools as canceled
252+ if processErr := a.processEvent(ctx, sessionID, &assistantMsg, event); processErr != nil {
257− for i, toolCall := range toolCalls {
253+ a.finishMessage(ctx, &assistantMsg, message.FinishReasonCanceled)
258− toolResults[i] = message.ToolResult{
254+ return assistantMsg, nil, processErr
259− ToolCallID: toolCall.ID,
260− Content: "Tool execution canceled by user",
256+ if ctx.Err() != nil {
257+ a.finishMessage(context.Background(), &assistantMsg, message.FinishReasonCanceled)
258+ return assistantMsg, nil, ctx.Err()
264− return toolResults, ctx.Err()
262+ toolResults := make([]message.ToolResult, len(assistantMsg.ToolCalls()))
263+ toolCalls := assistantMsg.ToolCalls()
267 for i, toolCall := range toolCalls {
264 for i, toolCall := range toolCalls {
268− // Check for cancellation before executing each tool
270 case <-ctx.Done():
266 case <-ctx.Done():
271− // Mark this and all remaining tools as canceled
267+ a.finishMessage(context.Background(), &assistantMsg, message.FinishReasonCanceled)
268+ // Make all future tool calls cancelled
272 for j := i; j < len(toolCalls); j++ {
269 for j := i; j < len(toolCalls); j++ {
273 toolResults[j] = message.ToolResult{
270 toolResults[j] = message.ToolResult{
274 ToolCallID: toolCalls[j].ID,
271 ToolCallID: toolCalls[j].ID,
@@ -276,412 +273,180 @@ func (a *agent) ExecuteTools(ctx context.Context, toolCalls []message.ToolCall,
276 IsError: true,
273 IsError: true,
279− return toolResults, ctx.Err()
276+ goto out
281 // Continue processing
278 // Continue processing
282− }
279+ var tool tools.BaseTool
288− // Find and execute the appropriate tool
289− for _, tool := range tls {
280+ for _, availableTools := range a.tools {
290− if tool.Info().Name == toolCall.Name {
281+ if availableTools.Info().Name == toolCall.Name {
291− found = true
282+ tool = availableTools
292− toolResult, toolErr := tool.Run(ctx, tools.ToolCall{
294− Name: toolCall.Name,
295− Input: toolCall.Input,
299− if errors.Is(toolErr, context.Canceled) {
300− response = "Tool execution canceled by user"
302− response = fmt.Sprintf("Error running tool: %s", toolErr)
306− response = toolResult.Content
307− isError = toolResult.IsError
314− response = fmt.Sprintf("Tool not found: %s", toolCall.Name)
318− toolResults[i] = message.ToolResult{
319− ToolCallID: toolCall.ID,
325− return toolResults, nil
286+ // Tool not found
326−}
287+ if tool == nil {
328−// handleToolExecution processes tool calls and creates tool result messages
329−func (a *agent) handleToolExecution(
330− ctx context.Context,
331− assistantMsg message.Message,
332−) (*message.Message, error) {
335− // If cancelled, create tool results that indicate cancellation
336− if len(assistantMsg.ToolCalls()) > 0 {
337− toolResults := make([]message.ToolResult, 0, len(assistantMsg.ToolCalls()))
338− for _, tc := range assistantMsg.ToolCalls() {
339− toolResults = append(toolResults, message.ToolResult{
288+ toolResults[i] = message.ToolResult{
340− ToolCallID: tc.ID,
289+ ToolCallID: toolCall.ID,
341− Content: "Tool execution canceled by user",
290+ Content: fmt.Sprintf("Tool not found: %s", toolCall.Name),
342 IsError: true,
291 IsError: true,
346− // Use background context to ensure the message is created even if original context is cancelled
296+ toolResult, toolErr := tool.Run(ctx, tools.ToolCall{
347− bgCtx := context.Background()
297+ ID: toolCall.ID,
348− parts := make([]message.ContentPart, 0)
298+ Name: toolCall.Name,
349− for _, toolResult := range toolResults {
299+ Input: toolCall.Input,
350− parts = append(parts, toolResult)
352− msg, err := a.messages.Create(bgCtx, assistantMsg.SessionID, message.CreateMessageParams{
356− if err != nil {
301+ if toolErr != nil {
357− return nil, fmt.Errorf("failed to create cancelled tool message: %w", err)
302+ if errors.Is(toolErr, permission.ErrorPermissionDenied) {
358− }
303+ toolResults[i] = message.ToolResult{
359− return &msg, ctx.Err()
304+ ToolCallID: toolCall.ID,
305+ Content: "Permission denied",
308+ for j := i + 1; j < len(toolCalls); j++ {
309+ toolResults[j] = message.ToolResult{
310+ ToolCallID: toolCalls[j].ID,
311+ Content: "Tool execution canceled by user",
315+ a.finishMessage(ctx, &assistantMsg, message.FinishReasonPermissionDenied)
317+ toolResults[i] = message.ToolResult{
318+ ToolCallID: toolCall.ID,
319+ Content: toolErr.Error(),
322+ for j := i; j < len(toolCalls); j++ {
323+ toolResults[j] = message.ToolResult{
324+ ToolCallID: toolCalls[j].ID,
325+ Content: "Previous tool failed",
361− return nil, ctx.Err()
363− // Continue processing
365−
329+ a.finishMessage(ctx, &assistantMsg, message.FinishReasonError)
366− if len(assistantMsg.ToolCalls()) == 0 {
330+ }
367− return nil, nil
331+ // If permission is denied or an error happens we cancel all the following tools
370− toolResults, err := a.ExecuteTools(ctx, assistantMsg.ToolCalls(), a.tools)
372− // If error is from cancellation, still return the partial results we have
373− if errors.Is(err, context.Canceled) {
374− // Use background context to ensure the message is created even if original context is cancelled
375− bgCtx := context.Background()
376− parts := make([]message.ContentPart, 0)
377− for _, toolResult := range toolResults {
378− parts = append(parts, toolResult)
380−
334+ toolResults[i] = message.ToolResult{
381− msg, createErr := a.messages.Create(bgCtx, assistantMsg.SessionID, message.CreateMessageParams{
335+ ToolCallID: toolCall.ID,
382− Role: message.Tool,
336+ Content: toolResult.Content,
383− Parts: parts,
337+ Metadata: toolResult.Metadata,
384− })
338+ IsError: toolResult.IsError,
385− if createErr != nil {
386− logging.ErrorPersist(fmt.Sprintf("Failed to create tool message after cancellation: %v", createErr))
394− parts := make([]message.ContentPart, 0, len(toolResults))
343+ if len(toolResults) == 0 {
395− for _, toolResult := range toolResults {
344+ return assistantMsg, nil, nil
396− parts = append(parts, toolResult)
398−
346+ parts := make([]message.ContentPart, 0)
347+ for _, tr := range toolResults {
348+ parts = append(parts, tr)
399− msg, err := a.messages.Create(ctx, assistantMsg.SessionID, message.CreateMessageParams{
350+ msg, err := a.messages.Create(context.Background(), assistantMsg.SessionID, message.CreateMessageParams{
400 Role: message.Tool,
351 Role: message.Tool,
401 Parts: parts,
352 Parts: parts,
403 if err != nil {
354 if err != nil {
404− return nil, fmt.Errorf("failed to create tool message: %w", err)
355+ return assistantMsg, nil, fmt.Errorf("failed to create cancelled tool message: %w", err)
407− return &msg, nil
358+ return assistantMsg, &msg, err
410−// generate handles the main generation workflow
361+func (a *agent) finishMessage(ctx context.Context, msg *message.Message, finishReson message.FinishReason) {
411−func (a *agent) generate(ctx context.Context, sessionID string, content string) error {
362+ msg.AddFinish(finishReson)
412− ctx = context.WithValue(ctx, tools.SessionIDContextKey, sessionID)
363+ _ = a.messages.Update(ctx, *msg)
414− // Handle context cancellation at any point
366+func (a *agent) processEvent(ctx context.Context, sessionID string, assistantMsg *message.Message, event provider.ProviderEvent) error {
415− if err := ctx.Err(); err != nil {
369+ return ctx.Err()
416− return ErrRequestCancelled
370+ default:
371+ // Continue processing.
419− messages, err := a.messages.List(ctx, sessionID)
374+ switch event.Type {
420− if err != nil {
375+ case provider.EventThinkingDelta:
376+ assistantMsg.AppendReasoningContent(event.Content)
377+ return a.messages.Update(ctx, *assistantMsg)
378+ case provider.EventContentDelta:
379+ assistantMsg.AppendContent(event.Content)
380+ return a.messages.Update(ctx, *assistantMsg)
381+ case provider.EventError:
382+ if errors.Is(event.Error, context.Canceled) {
383+ logging.InfoPersist(fmt.Sprintf("Event processing canceled for session: %s", sessionID))
384+ return context.Canceled
386+ logging.ErrorPersist(event.Error.Error())
388+ case provider.EventComplete:
389+ assistantMsg.SetToolCalls(event.Response.ToolCalls)
390+ assistantMsg.AddFinish(event.Response.FinishReason)
391+ if err := a.messages.Update(ctx, *assistantMsg); err != nil {
421− return fmt.Errorf("failed to list messages: %w", err)
392+ return fmt.Errorf("failed to update message: %w", err)
394+ return a.TrackUsage(ctx, sessionID, a.provider.Model(), event.Response.Usage)
424− if len(messages) == 0 {
397+ return nil
425− titleCtx := context.Background()
426− go a.handleTitleGeneration(titleCtx, sessionID, content)
429− userMsg, err := a.messages.Create(ctx, sessionID, message.CreateMessageParams{
400+func (a *agent) TrackUsage(ctx context.Context, sessionID string, model models.Model, usage provider.TokenUsage) error {
430− Role: message.User,
401+ sess, err := a.sessions.Get(ctx, sessionID)
431− Parts: []message.ContentPart{
432− message.TextContent{
437 if err != nil {
402 if err != nil {
438− return fmt.Errorf("failed to create user message: %w", err)
403+ return fmt.Errorf("failed to get session: %w", err)
441− messages = append(messages, userMsg)
406+ cost := model.CostPer1MInCached/1e6*float64(usage.CacheCreationTokens) +
442−
407+ model.CostPer1MOutCached/1e6*float64(usage.CacheReadTokens) +
443− for {
408+ model.CostPer1MIn/1e6*float64(usage.InputTokens) +
444− // Check for cancellation before each iteration
409+ model.CostPer1MOut/1e6*float64(usage.OutputTokens)
447− return ErrRequestCancelled
449− // Continue processing
452− eventChan, err := a.agent.StreamResponse(ctx, messages, a.tools)
454− if errors.Is(err, context.Canceled) {
455− return ErrRequestCancelled
457− return fmt.Errorf("failed to stream response: %w", err)
460− assistantMsg, err := a.messages.Create(ctx, sessionID, message.CreateMessageParams{
461− Role: message.Assistant,
462− Parts: []message.ContentPart{},
466− return fmt.Errorf("failed to create assistant message: %w", err)
469− ctx = context.WithValue(ctx, tools.MessageIDContextKey, assistantMsg.ID)
471− // Process events from the LLM provider
472− for event := range eventChan {
473− if err := a.processEvent(ctx, sessionID, &assistantMsg, event); err != nil {
474− if errors.Is(err, context.Canceled) {
475− // Mark as canceled but don't create separate message
476− assistantMsg.AddFinish("canceled")
477− _ = a.messages.Update(context.Background(), assistantMsg)
478− return ErrRequestCancelled
480− assistantMsg.AddFinish("error:" + err.Error())
481− _ = a.messages.Update(ctx, assistantMsg)
482− return fmt.Errorf("event processing error: %w", err)
485− // Check for cancellation during event processing
489− assistantMsg.AddFinish("canceled")
490− _ = a.messages.Update(context.Background(), assistantMsg)
491− return ErrRequestCancelled
496− // Check for cancellation before tool execution
499− assistantMsg.AddFinish("canceled_by_user")
500− _ = a.messages.Update(context.Background(), assistantMsg)
501− return ErrRequestCancelled
505− // Execute any tool calls
506− toolMsg, err := a.handleToolExecution(ctx, assistantMsg)
508− if errors.Is(err, context.Canceled) {
509− assistantMsg.AddFinish("canceled_by_user")
510− _ = a.messages.Update(context.Background(), assistantMsg)
511− return ErrRequestCancelled
513− return fmt.Errorf("tool execution error: %w", err)
516− if err := a.messages.Update(ctx, assistantMsg); err != nil {
517− return fmt.Errorf("failed to update assistant message: %w", err)
520− // If no tool calls, we're done
521− if len(assistantMsg.ToolCalls()) == 0 {
525− // Add messages for next iteration
411+ sess.Cost += cost
526− messages = append(messages, assistantMsg)
412+ sess.CompletionTokens += usage.OutputTokens
527− if toolMsg != nil {
413+ sess.PromptTokens += usage.InputTokens
528− messages = append(messages, *toolMsg)
531− // Check for cancellation after tool execution
415+ _, err = a.sessions.Save(ctx, sess)
532− select {
416+ if err != nil {
533− case <-ctx.Done():
417+ return fmt.Errorf("failed to save session: %w", err)
534− return ErrRequestCancelled
539 return nil
419 return nil
542−// getAgentProviders initializes the LLM providers based on the chosen model
422+func createAgentProvider(agentName config.AgentName) (provider.Provider, error) {
543−func getAgentProviders(ctx context.Context, model models.Model) (provider.Provider, provider.Provider, error) {
423+ cfg := config.Get()
544− maxTokens := config.Get().Model.CoderMaxTokens
424+ agentConfig, ok := cfg.Agents[agentName]
546− providerConfig, ok := config.Get().Providers[model.Provider]
426+ return nil, fmt.Errorf("agent %s not found", agentName)
547− if !ok || providerConfig.Disabled {
427+ }
548− return nil, nil, ErrProviderNotEnabled
428+ model, ok := models.SupportedModels[agentConfig.Model]
430+ return nil, fmt.Errorf("model %s not supported", agentConfig.Model)
551− var agentProvider provider.Provider
433+ providerCfg, ok := cfg.Providers[model.Provider]
552− var titleGenerator provider.Provider
434+ if !ok {
553− var err error
435+ return nil, fmt.Errorf("provider %s not supported", model.Provider)
555− switch model.Provider {
437+ if providerCfg.Disabled {
556− case models.ProviderOpenAI:
438+ return nil, fmt.Errorf("provider %s is not enabled", model.Provider)
557− agentProvider, err = provider.NewOpenAIProvider(
558− provider.WithOpenAISystemMessage(
559− prompt.CoderOpenAISystemPrompt(),
561− provider.WithOpenAIMaxTokens(maxTokens),
562− provider.WithOpenAIModel(model),
563− provider.WithOpenAIKey(providerConfig.APIKey),
566− return nil, nil, fmt.Errorf("failed to create OpenAI agent provider: %w", err)
568−
440+ agentProvider, err := provider.NewProvider(
569− titleGenerator, err = provider.NewOpenAIProvider(
441+ model.Provider,
570− provider.WithOpenAISystemMessage(
442+ provider.WithAPIKey(providerCfg.APIKey),
571− prompt.TitlePrompt(),
443+ provider.WithModel(model),
572− ),
444+ provider.WithSystemMessage(prompt.GetAgentPrompt(agentName, model.Provider)),
573− provider.WithOpenAIMaxTokens(80),
445+ provider.WithMaxTokens(agentConfig.MaxTokens),
574− provider.WithOpenAIModel(model),
575− provider.WithOpenAIKey(providerConfig.APIKey),
577− if err != nil {
447+ if err != nil {
578− return nil, nil, fmt.Errorf("failed to create OpenAI title generator: %w", err)
448+ return nil, fmt.Errorf("could not create provider: %v", err)
581− case models.ProviderAnthropic:
582− agentProvider, err = provider.NewAnthropicProvider(
583− provider.WithAnthropicSystemMessage(
584− prompt.CoderAnthropicSystemPrompt(),
586− provider.WithAnthropicMaxTokens(maxTokens),
587− provider.WithAnthropicKey(providerConfig.APIKey),
588− provider.WithAnthropicModel(model),
591− return nil, nil, fmt.Errorf("failed to create Anthropic agent provider: %w", err)
594− titleGenerator, err = provider.NewAnthropicProvider(
595− provider.WithAnthropicSystemMessage(
596− prompt.TitlePrompt(),
598− provider.WithAnthropicMaxTokens(80),
599− provider.WithAnthropicKey(providerConfig.APIKey),
600− provider.WithAnthropicModel(model),
603− return nil, nil, fmt.Errorf("failed to create Anthropic title generator: %w", err)
606− case models.ProviderGemini:
607− agentProvider, err = provider.NewGeminiProvider(
609− provider.WithGeminiSystemMessage(
610− prompt.CoderOpenAISystemPrompt(),
612− provider.WithGeminiMaxTokens(int32(maxTokens)),
613− provider.WithGeminiKey(providerConfig.APIKey),
614− provider.WithGeminiModel(model),
617− return nil, nil, fmt.Errorf("failed to create Gemini agent provider: %w", err)
620− titleGenerator, err = provider.NewGeminiProvider(
622− provider.WithGeminiSystemMessage(
623− prompt.TitlePrompt(),
625− provider.WithGeminiMaxTokens(80),
626− provider.WithGeminiKey(providerConfig.APIKey),
627− provider.WithGeminiModel(model),
630− return nil, nil, fmt.Errorf("failed to create Gemini title generator: %w", err)
633− case models.ProviderGROQ:
634− agentProvider, err = provider.NewOpenAIProvider(
635− provider.WithOpenAISystemMessage(
636− prompt.CoderAnthropicSystemPrompt(),
638− provider.WithOpenAIMaxTokens(maxTokens),
639− provider.WithOpenAIModel(model),
640− provider.WithOpenAIKey(providerConfig.APIKey),
641− provider.WithOpenAIBaseURL("https://api.groq.com/openai/v1"),
644− return nil, nil, fmt.Errorf("failed to create GROQ agent provider: %w", err)
647− titleGenerator, err = provider.NewOpenAIProvider(
648− provider.WithOpenAISystemMessage(
649− prompt.TitlePrompt(),
651− provider.WithOpenAIMaxTokens(80),
652− provider.WithOpenAIModel(model),
653− provider.WithOpenAIKey(providerConfig.APIKey),
654− provider.WithOpenAIBaseURL("https://api.groq.com/openai/v1"),
657− return nil, nil, fmt.Errorf("failed to create GROQ title generator: %w", err)
660− case models.ProviderBedrock:
661− agentProvider, err = provider.NewBedrockProvider(
662− provider.WithBedrockSystemMessage(
663− prompt.CoderAnthropicSystemPrompt(),
665− provider.WithBedrockMaxTokens(maxTokens),
666− provider.WithBedrockModel(model),
669− return nil, nil, fmt.Errorf("failed to create Bedrock agent provider: %w", err)
672− titleGenerator, err = provider.NewBedrockProvider(
673− provider.WithBedrockSystemMessage(
674− prompt.TitlePrompt(),
676− provider.WithBedrockMaxTokens(80),
677− provider.WithBedrockModel(model),
680− return nil, nil, fmt.Errorf("failed to create Bedrock title generator: %w", err)
683− return nil, nil, fmt.Errorf("unsupported provider: %s", model.Provider)
686− return agentProvider, titleGenerator, nil
451+ return agentProvider, nil