cdc5f20 · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

cleanup diff, cleanup agent

cdc5f20Kujtim Hoxha2 years ago

README.md+4−10

compare with main

unifiedsplit

@@ -53,6 +53,7 @@ termai -d

53### Keyboard Shortcuts

55#### Global Shortcuts

56+

56- `?`: Toggle help panel

57- `?`: Toggle help panel

57- `Ctrl+C` or `q`: Quit application

58- `Ctrl+C` or `q`: Quit application

58- `L`: View logs

59- `L`: View logs

@@ -60,10 +61,12 @@ termai -d

60- `Esc`: Close current view/dialog or return to normal mode

61- `Esc`: Close current view/dialog or return to normal mode

62#### Session Management

63#### Session Management

64+

63- `N`: Create new session

65- `N`: Create new session

64- `Enter` or `Space`: Select session (in sessions list)

66- `Enter` or `Space`: Select session (in sessions list)

66#### Editor Shortcuts (Vim-like)

68#### Editor Shortcuts (Vim-like)

69+

67- `i`: Enter insert mode

70- `i`: Enter insert mode

68- `Esc`: Enter normal mode

71- `Esc`: Enter normal mode

69- `v`: Enter visual mode

72- `v`: Enter visual mode

@@ -72,6 +75,7 @@ termai -d

72- `Ctrl+S`: Send message (in insert mode)

75- `Ctrl+S`: Send message (in insert mode)

74#### Navigation

77#### Navigation

78+

75- Arrow keys: Navigate through lists and content

79- Arrow keys: Navigate through lists and content

76- Page Up/Down: Scroll through content

80- Page Up/Down: Scroll through content

@@ -112,16 +116,6 @@ go build -o termai

112./termai

116./termai

113```

117```

114

118

115−### Important: Building the Diff Script

116−

117−Before building or running the application, you must first build the diff script by running:

118−

119−```bash

120−go run cmd/diff/main.go

121−```

122−

123−This command generates the necessary JavaScript file (`index.mjs`) used by the diff functionality in the application.

124−

125## Acknowledgments

119## Acknowledgments

126

120

127TermAI builds upon the work of several open source projects and developers:

121TermAI builds upon the work of several open source projects and developers:

cmd/diff/main.godeleted−102

@@ -1,102 +0,0 @@

1−package main

2−

3−import (

4− "fmt"

5− "io"

6− "os"

7− "os/exec"

8− "path/filepath"

9−)

10−

11−func main() {

12− // Create a temporary directory

13− tempDir, err := os.MkdirTemp("", "git-split-diffs")

14− if err != nil {

15− fmt.Printf("Error creating temp directory: %v\n", err)

16− os.Exit(1)

17− }

18− defer func() {

19− fmt.Printf("Cleaning up temporary directory: %s\n", tempDir)

20− os.RemoveAll(tempDir)

21− }()

22− fmt.Printf("Created temporary directory: %s\n", tempDir)

23−

24− // Clone the repository with minimum depth

25− fmt.Println("Cloning git-split-diffs repository with minimum depth...")

26− cmd := exec.Command("git", "clone", "--depth=1", "https://github.com/kujtimiihoxha/git-split-diffs", tempDir)

27− cmd.Stdout = os.Stdout

28− cmd.Stderr = os.Stderr

29− if err := cmd.Run(); err != nil {

30− fmt.Printf("Error cloning repository: %v\n", err)

31− os.Exit(1)

32− }

33−

34− // Run npm install

35− fmt.Println("Running npm install...")

36− cmdNpmInstall := exec.Command("npm", "install")

37− cmdNpmInstall.Dir = tempDir

38− cmdNpmInstall.Stdout = os.Stdout

39− cmdNpmInstall.Stderr = os.Stderr

40− if err := cmdNpmInstall.Run(); err != nil {

41− fmt.Printf("Error running npm install: %v\n", err)

42− os.Exit(1)

43− }

44−

45− // Run npm run build

46− fmt.Println("Running npm run build...")

47− cmdNpmBuild := exec.Command("npm", "run", "build")

48− cmdNpmBuild.Dir = tempDir

49− cmdNpmBuild.Stdout = os.Stdout

50− cmdNpmBuild.Stderr = os.Stderr

51− if err := cmdNpmBuild.Run(); err != nil {

52− fmt.Printf("Error running npm run build: %v\n", err)

53− os.Exit(1)

54− }

55−

56− destDir := filepath.Join(".", "internal", "assets", "diff")

57− destFile := filepath.Join(destDir, "index.mjs")

58−

59− // Make sure the destination directory exists

60− if err := os.MkdirAll(destDir, 0o755); err != nil {

61− fmt.Printf("Error creating destination directory: %v\n", err)

62− os.Exit(1)

63− }

64−

65− // Copy the file

66− srcFile := filepath.Join(tempDir, "build", "index.mjs")

67− fmt.Printf("Copying %s to %s\n", srcFile, destFile)

68− if err := copyFile(srcFile, destFile); err != nil {

69− fmt.Printf("Error copying file: %v\n", err)

70− os.Exit(1)

71− }

72−

73− fmt.Println("Successfully completed the process!")

74−}

75−

76−// copyFile copies a file from src to dst

77−func copyFile(src, dst string) error {

78− sourceFile, err := os.Open(src)

79− if err != nil {

80− return err

81− }

82− defer sourceFile.Close()

83−

84− destFile, err := os.Create(dst)

85− if err != nil {

86− return err

87− }

88− defer destFile.Close()

89−

90− _, err = io.Copy(destFile, sourceFile)

91− if err != nil {

92− return err

93− }

94−

95− // Make sure the file is written to disk

96− err = destFile.Sync()

97− if err != nil {

98− return err

99− }

100−

101− return nil

102−}

cmd/root.go+5−7

@@ -9,7 +9,6 @@ import (

10 tea "github.com/charmbracelet/bubbletea"

11 "github.com/kujtimiihoxha/termai/internal/app"

12− "github.com/kujtimiihoxha/termai/internal/assets"

13 "github.com/kujtimiihoxha/termai/internal/config"

12 "github.com/kujtimiihoxha/termai/internal/config"

14 "github.com/kujtimiihoxha/termai/internal/db"

13 "github.com/kujtimiihoxha/termai/internal/db"

15 "github.com/kujtimiihoxha/termai/internal/llm/agent"

14 "github.com/kujtimiihoxha/termai/internal/llm/agent"

@@ -52,11 +51,6 @@ var rootCmd = &cobra.Command{

52 return err

51 return err

53 }

52 }

55− err = assets.WriteAssets()

56− if err != nil {

57− logging.Error("Error writing assets: %v", err)

58− }

59−

60 // Connect DB, this will also run migrations

54 // Connect DB, this will also run migrations

61 conn, err := db.Connect()

55 conn, err := db.Connect()

62 if err != nil {

56 if err != nil {

@@ -67,7 +61,11 @@ var rootCmd = &cobra.Command{

67 ctx, cancel := context.WithCancel(context.Background())

61 ctx, cancel := context.WithCancel(context.Background())

68 defer cancel()

62 defer cancel()

70− app := app.New(ctx, conn)

64+ app, err := app.New(ctx, conn)

65+ if err != nil {

66+ logging.Error("Failed to create app: %v", err)

67+ return err

68+ }

72 // Set up the TUI

70 // Set up the TUI

73 zone.NewGlobal()

71 zone.NewGlobal()

go.mod+1−1

@@ -31,7 +31,6 @@ require (

31 github.com/muesli/reflow v0.3.0

32 github.com/muesli/termenv v0.16.0

33 github.com/openai/openai-go v0.1.0-beta.2

34− github.com/sergi/go-diff v1.3.2-0.20230802210424-5b0b94c5c0d3

35 github.com/spf13/cobra v1.9.1

34 github.com/spf13/cobra v1.9.1

36 github.com/spf13/viper v1.20.0

35 github.com/spf13/viper v1.20.0

37 github.com/stretchr/testify v1.10.0

36 github.com/stretchr/testify v1.10.0

@@ -107,6 +106,7 @@ require (

107 github.com/rivo/uniseg v0.4.7 // indirect

106 github.com/rivo/uniseg v0.4.7 // indirect

108 github.com/sagikazarmark/locafero v0.7.0 // indirect

107 github.com/sagikazarmark/locafero v0.7.0 // indirect

109 github.com/sahilm/fuzzy v0.1.1 // indirect

108 github.com/sahilm/fuzzy v0.1.1 // indirect

109+ github.com/sergi/go-diff v1.3.2-0.20230802210424-5b0b94c5c0d3 // indirect

110 github.com/skeema/knownhosts v1.3.1 // indirect

111 github.com/sourcegraph/conc v0.3.0 // indirect

112 github.com/spf13/afero v1.12.0 // indirect

internal/app/app.go+18−2

@@ -9,6 +9,7 @@ import (

10 "github.com/kujtimiihoxha/termai/internal/db"

11 "github.com/kujtimiihoxha/termai/internal/history"

12+ "github.com/kujtimiihoxha/termai/internal/llm/agent"

12 "github.com/kujtimiihoxha/termai/internal/logging"

13 "github.com/kujtimiihoxha/termai/internal/logging"

13 "github.com/kujtimiihoxha/termai/internal/lsp"

14 "github.com/kujtimiihoxha/termai/internal/lsp"

14 "github.com/kujtimiihoxha/termai/internal/message"

15 "github.com/kujtimiihoxha/termai/internal/message"

@@ -22,6 +23,8 @@ type App struct {

22 Files history.Service

23 Files history.Service

23 Permissions permission.Service

24 Permissions permission.Service

26+ CoderAgent agent.Service

27+

25 LSPClients map[string]*lsp.Client

28 LSPClients map[string]*lsp.Client

27 clientsMutex sync.RWMutex

30 clientsMutex sync.RWMutex

@@ -31,7 +34,7 @@ type App struct {

31 watcherWG sync.WaitGroup

34 watcherWG sync.WaitGroup

32}

35}

34−func New(ctx context.Context, conn *sql.DB) *App {

37+func New(ctx context.Context, conn *sql.DB) (*App, error) {

35 q := db.New(conn)

38 q := db.New(conn)

36 sessions := session.NewService(q)

39 sessions := session.NewService(q)

37 messages := message.NewService(q)

40 messages := message.NewService(q)

@@ -45,9 +48,22 @@ func New(ctx context.Context, conn *sql.DB) *App {

45 LSPClients: make(map[string]*lsp.Client),

48 LSPClients: make(map[string]*lsp.Client),

46 }

49 }

51+ var err error

52+ app.CoderAgent, err = agent.NewCoderAgent(

53+

54+ app.Permissions,

55+ app.Sessions,

56+ app.Messages,

57+ app.LSPClients,

58+ )

59+ if err != nil {

60+ logging.Error("Failed to create coder agent", err)

61+ return nil, err

62+ }

63+

48 app.initLSPClients(ctx)

64 app.initLSPClients(ctx)

50− return app

66+ return app, nil

51}

67}

53// Shutdown performs a clean shutdown of the application

69// Shutdown performs a clean shutdown of the application

internal/assets/diff/themes/dark.jsondeleted−73

@@ -1,73 +0,0 @@

1−{

2− "SYNTAX_HIGHLIGHTING_THEME": "dark-plus",

3− "DEFAULT_COLOR": {

4− "color": "#ffffff",

5− "backgroundColor": "#212121"

6− },

7− "COMMIT_HEADER_COLOR": {

8− "color": "#cccccc"

9− },

10− "COMMIT_HEADER_LABEL_COLOR": {

11− "color": "#00000022"

12− },

13− "COMMIT_SHA_COLOR": {

14− "color": "#00eeaa"

15− },

16− "COMMIT_AUTHOR_COLOR": {

17− "color": "#00aaee"

18− },

19− "COMMIT_DATE_COLOR": {

20− "color": "#cccccc"

21− },

22− "COMMIT_MESSAGE_COLOR": {

23− "color": "#cccccc"

24− },

25− "COMMIT_TITLE_COLOR": {

26− "modifiers": [

27− "bold"

28− ]

29− },

30− "FILE_NAME_COLOR": {

31− "color": "#ffdd99"

32− },

33− "BORDER_COLOR": {

34− "color": "#ffdd9966",

35− "modifiers": [

36− "dim"

37− ]

38− },

39− "HUNK_HEADER_COLOR": {

40− "modifiers": [

41− "dim"

42− ]

43− },

44− "DELETED_WORD_COLOR": {

45− "color": "#ffcccc",

46− "backgroundColor": "#ff000033"

47− },

48− "INSERTED_WORD_COLOR": {

49− "color": "#ccffcc",

50− "backgroundColor": "#00ff0033"

51− },

52− "DELETED_LINE_NO_COLOR": {

53− "color": "#00000022",

54− "backgroundColor": "#00000022"

55− },

56− "INSERTED_LINE_NO_COLOR": {

57− "color": "#00000022",

58− "backgroundColor": "#00000022"

59− },

60− "UNMODIFIED_LINE_NO_COLOR": {

61− "color": "#666666"

62− },

63− "DELETED_LINE_COLOR": {

64− "color": "#cc6666",

65− "backgroundColor": "#3a3030"

66− },

67− "INSERTED_LINE_COLOR": {

68− "color": "#66cc66",

69− "backgroundColor": "#303a30"

70− },

71− "UNMODIFIED_LINE_COLOR": {},

72− "MISSING_LINE_COLOR": {}

73−}

internal/assets/embed.godeleted−6

@@ -1,6 +0,0 @@

1−package assets

2−

3−import "embed"

4−

5−//go:embed diff

6−var FS embed.FS

internal/assets/write.godeleted−60

@@ -1,60 +0,0 @@

1−package assets

2−

3−import (

4− "os"

5− "path/filepath"

6−

7− "github.com/kujtimiihoxha/termai/internal/config"

8−)

9−

10−func WriteAssets() error {

11− appCfg := config.Get()

12− appWd := config.WorkingDirectory()

13− scriptDir := filepath.Join(

14− appWd,

15− appCfg.Data.Directory,

16− "diff",

17− )

18− scriptPath := filepath.Join(scriptDir, "index.mjs")

19− // Before, run the script in cmd/diff/main.go to build this file

20− if _, err := os.Stat(scriptPath); err != nil {

21− scriptData, err := FS.ReadFile("diff/index.mjs")

22− if err != nil {

23− return err

24− }

25−

26− err = os.MkdirAll(scriptDir, 0o755)

27− if err != nil {

28− return err

29− }

30− err = os.WriteFile(scriptPath, scriptData, 0o755)

31− if err != nil {

32− return err

33− }

34− }

35−

36− themeDir := filepath.Join(

37− appWd,

38− appCfg.Data.Directory,

39− "themes",

40− )

41−

42− themePath := filepath.Join(themeDir, "dark.json")

43−

44− if _, err := os.Stat(themePath); err != nil {

45− themeData, err := FS.ReadFile("diff/themes/dark.json")

46− if err != nil {

47− return err

48− }

49−

50− err = os.MkdirAll(themeDir, 0o755)

51− if err != nil {

52− return err

53− }

54− err = os.WriteFile(themePath, themeData, 0o755)

55− if err != nil {

56− return err

57− }

58− }

59− return nil

60−}

internal/git/diff.go+17−18

@@ -11,7 +11,6 @@ import (

12 "github.com/go-git/go-git/v5"

13 "github.com/go-git/go-git/v5/plumbing/object"

14− "github.com/kujtimiihoxha/termai/internal/config"

15)

14)

17type DiffStats struct {

16type DiffStats struct {

@@ -197,32 +196,32 @@ func isSplitDiffsAvailable() bool {

197}

196}

198

197

199func formatWithSplitDiffs(diffText string, width int) (string, error) {

198func formatWithSplitDiffs(diffText string, width int) (string, error) {

200− var cmd *exec.Cmd

199+ args := []string{

200+ "--color",

201+ }

201

202

202− appCfg := config.Get()

203+ var diffCmd *exec.Cmd

203− appWd := config.WorkingDirectory()

204− script := filepath.Join(

205− appWd,

206− appCfg.Data.Directory,

207− "diff",

208− "index.mjs",

209− )

210

204

211− cmd = exec.Command("node", script, "--color")

205+ if _, err := exec.LookPath("git-split-diffs-opencode"); err == nil {

206+ fullArgs := append([]string{"git-split-diffs-opencode"}, args...)

207+ diffCmd = exec.Command(fullArgs[0], fullArgs[1:]...)

208+ } else {

209+ npxArgs := append([]string{"git-split-diffs-opencode"}, args...)

210+ diffCmd = exec.Command("npx", npxArgs...)

211+ }

212

213− cmd.Env = append(os.Environ(), fmt.Sprintf("COLUMNS=%d", width))

213+ diffCmd.Env = append(os.Environ(), fmt.Sprintf("DIFF_COLUMNS=%d", width))

214

215− cmd.Stdin = strings.NewReader(diffText)

215+ diffCmd.Stdin = strings.NewReader(diffText)

216

217 var out bytes.Buffer

218− cmd.Stdout = &out

218+ diffCmd.Stdout = &out

219

220 var stderr bytes.Buffer

221− cmd.Stderr = &stderr

221+ diffCmd.Stderr = &stderr

222

223− err := cmd.Run()

224− if err != nil {

223+ if err := diffCmd.Run(); err != nil {

225− return "", fmt.Errorf("git-split-diffs error: %v, stderr: %s", err, stderr.String())

224+ return "", fmt.Errorf("git-split-diffs-opencode error: %w, stderr: %s", err, stderr.String())

226 }

225 }

227

226

228 return out.String(), nil

227 return out.String(), nil

internal/llm/agent/agent-tool.go+22−12

@@ -5,14 +5,16 @@ import (

5 "encoding/json"

6 "fmt"

8− "github.com/kujtimiihoxha/termai/internal/app"

9 "github.com/kujtimiihoxha/termai/internal/llm/tools"

8 "github.com/kujtimiihoxha/termai/internal/llm/tools"

9+ "github.com/kujtimiihoxha/termai/internal/lsp"

10 "github.com/kujtimiihoxha/termai/internal/message"

11+ "github.com/kujtimiihoxha/termai/internal/session"

11)

12)

13type agentTool struct {

14type agentTool struct {

14− parentSessionID string

15+ sessions session.Service

15− app *app.App

16+ messages message.Service

17+ lspClients map[string]*lsp.Client

16}

18}

18const (

20const (

@@ -46,12 +48,17 @@ func (b *agentTool) Run(ctx context.Context, call tools.ToolCall) (tools.ToolRes

46 return tools.NewTextErrorResponse("prompt is required"), nil

48 return tools.NewTextErrorResponse("prompt is required"), nil

47 }

49 }

51+ sessionID, messageID := tools.GetContextValues(ctx)

52+ if sessionID == "" || messageID == "" {

53+ return tools.NewTextErrorResponse("session ID and message ID are required"), nil

54+ }

55+

49− agent, err := NewTaskAgent(b.app)

56+ agent, err := NewTaskAgent(b.lspClients)

50 if err != nil {

57 if err != nil {

51 return tools.NewTextErrorResponse(fmt.Sprintf("error creating agent: %s", err)), nil

58 return tools.NewTextErrorResponse(fmt.Sprintf("error creating agent: %s", err)), nil

52 }

59 }

54− session, err := b.app.Sessions.CreateTaskSession(ctx, call.ID, b.parentSessionID, "New Agent Session")

61+ session, err := b.sessions.CreateTaskSession(ctx, call.ID, sessionID, "New Agent Session")

55 if err != nil {

62 if err != nil {

56 return tools.NewTextErrorResponse(fmt.Sprintf("error creating session: %s", err)), nil

63 return tools.NewTextErrorResponse(fmt.Sprintf("error creating session: %s", err)), nil

57 }

64 }

@@ -61,7 +68,7 @@ func (b *agentTool) Run(ctx context.Context, call tools.ToolCall) (tools.ToolRes

61 return tools.NewTextErrorResponse(fmt.Sprintf("error generating agent: %s", err)), nil

68 return tools.NewTextErrorResponse(fmt.Sprintf("error generating agent: %s", err)), nil

62 }

69 }

64− messages, err := b.app.Messages.List(ctx, session.ID)

71+ messages, err := b.messages.List(ctx, session.ID)

65 if err != nil {

72 if err != nil {

66 return tools.NewTextErrorResponse(fmt.Sprintf("error listing messages: %s", err)), nil

73 return tools.NewTextErrorResponse(fmt.Sprintf("error listing messages: %s", err)), nil

67 }

74 }

@@ -74,11 +81,11 @@ func (b *agentTool) Run(ctx context.Context, call tools.ToolCall) (tools.ToolRes

74 return tools.NewTextErrorResponse("no assistant message found"), nil

81 return tools.NewTextErrorResponse("no assistant message found"), nil

75 }

82 }

77− updatedSession, err := b.app.Sessions.Get(ctx, session.ID)

84+ updatedSession, err := b.sessions.Get(ctx, session.ID)

78 if err != nil {

85 if err != nil {

79 return tools.NewTextErrorResponse(fmt.Sprintf("error: %s", err)), nil

86 return tools.NewTextErrorResponse(fmt.Sprintf("error: %s", err)), nil

80 }

87 }

81− parentSession, err := b.app.Sessions.Get(ctx, b.parentSessionID)

88+ parentSession, err := b.sessions.Get(ctx, sessionID)

82 if err != nil {

89 if err != nil {

83 return tools.NewTextErrorResponse(fmt.Sprintf("error: %s", err)), nil

90 return tools.NewTextErrorResponse(fmt.Sprintf("error: %s", err)), nil

84 }

91 }

@@ -87,16 +94,19 @@ func (b *agentTool) Run(ctx context.Context, call tools.ToolCall) (tools.ToolRes

87 parentSession.PromptTokens += updatedSession.PromptTokens

94 parentSession.PromptTokens += updatedSession.PromptTokens

88 parentSession.CompletionTokens += updatedSession.CompletionTokens

95 parentSession.CompletionTokens += updatedSession.CompletionTokens

90− _, err = b.app.Sessions.Save(ctx, parentSession)

97+ _, err = b.sessions.Save(ctx, parentSession)

91 if err != nil {

98 if err != nil {

92 return tools.NewTextErrorResponse(fmt.Sprintf("error: %s", err)), nil

99 return tools.NewTextErrorResponse(fmt.Sprintf("error: %s", err)), nil

93 }

100 }

94 return tools.NewTextResponse(response.Content().String()), nil

101 return tools.NewTextResponse(response.Content().String()), nil

95}

102}

103

97−func NewAgentTool(parentSessionID string, app *app.App) tools.BaseTool {

104+func NewAgentTool(

105+ Sessions session.Service,

106+ Messages message.Service,

107+) tools.BaseTool {

98 return &agentTool{

108 return &agentTool{

99− parentSessionID: parentSessionID,

109+ sessions: Sessions,

100− app: app,

110+ messages: Messages,

101 }

111 }

102}

112}

internal/llm/agent/agent.go+335−187

@@ -7,7 +7,6 @@ import (

7 "strings"

8 "sync"

10− "github.com/kujtimiihoxha/termai/internal/app"

11 "github.com/kujtimiihoxha/termai/internal/config"

10 "github.com/kujtimiihoxha/termai/internal/config"

12 "github.com/kujtimiihoxha/termai/internal/llm/models"

11 "github.com/kujtimiihoxha/termai/internal/llm/models"

13 "github.com/kujtimiihoxha/termai/internal/llm/prompt"

12 "github.com/kujtimiihoxha/termai/internal/llm/prompt"

@@ -15,22 +14,118 @@ import (

15 "github.com/kujtimiihoxha/termai/internal/llm/tools"

14 "github.com/kujtimiihoxha/termai/internal/llm/tools"

16 "github.com/kujtimiihoxha/termai/internal/logging"

15 "github.com/kujtimiihoxha/termai/internal/logging"

17 "github.com/kujtimiihoxha/termai/internal/message"

16 "github.com/kujtimiihoxha/termai/internal/message"

17+ "github.com/kujtimiihoxha/termai/internal/session"

18)

20+// Common errors

21+var (

22+ ErrProviderNotEnabled = errors.New("provider is not enabled")

23+ ErrRequestCancelled = errors.New("request cancelled by user")

24+ ErrSessionBusy = errors.New("session is currently processing another request")

25+)

26+

27+// Service defines the interface for generating responses

20−type Agent interface {

28+type Service interface {

21 Generate(ctx context.Context, sessionID string, content string) error

29 Generate(ctx context.Context, sessionID string, content string) error

30+ Cancel(sessionID string) error

22}

31}

24type agent struct {

33type agent struct {

25− *app.App

34+ sessions session.Service

35+ messages message.Service

26 model models.Model

36 model models.Model

27 tools []tools.BaseTool

37 tools []tools.BaseTool

28 agent provider.Provider

38 agent provider.Provider

29 titleGenerator provider.Provider

39 titleGenerator provider.Provider

40+ activeRequests sync.Map // map[sessionID]context.CancelFunc

41+}

42+

43+// NewAgent creates a new agent instance with the given model and tools

44+func NewAgent(ctx context.Context, sessions session.Service, messages message.Service, model models.Model, tools []tools.BaseTool) (Service, error) {

45+ agentProvider, titleGenerator, err := getAgentProviders(ctx, model)

46+ if err != nil {

47+ return nil, fmt.Errorf("failed to initialize providers: %w", err)

48+ }

49+

50+ return &agent{

51+ model: model,

52+ tools: tools,

53+ sessions: sessions,

54+ messages: messages,

55+ agent: agentProvider,

56+ titleGenerator: titleGenerator,

57+ activeRequests: sync.Map{},

58+ }, nil

59+}

60+

61+// Cancel cancels an active request by session ID

62+func (a *agent) Cancel(sessionID string) error {

63+ if cancelFunc, exists := a.activeRequests.LoadAndDelete(sessionID); exists {

64+ if cancel, ok := cancelFunc.(context.CancelFunc); ok {

65+ logging.InfoPersist(fmt.Sprintf("Request cancellation initiated for session: %s", sessionID))

66+ cancel()

67+ return nil

68+ }

69+ }

70+ return errors.New("no active request found for this session")

30}

71}

73+// Generate starts the generation process

74+func (a *agent) Generate(ctx context.Context, sessionID string, content string) error {

75+ // Check if this session already has an active request

76+ if _, busy := a.activeRequests.Load(sessionID); busy {

77+ return ErrSessionBusy

78+ }

79+

80+ // Create a cancellable context

81+ genCtx, cancel := context.WithCancel(ctx)

82+

83+ // Store cancel function to allow user cancellation

84+ a.activeRequests.Store(sessionID, cancel)

85+

86+ // Launch the generation in a goroutine

87+ go func() {

88+ defer func() {

89+ if r := recover(); r != nil {

90+ logging.ErrorPersist(fmt.Sprintf("Panic in Generate: %v", r))

91+ }

92+ }()

93+ defer a.activeRequests.Delete(sessionID)

94+ defer cancel()

95+

96+ if err := a.generate(genCtx, sessionID, content); err != nil {

97+ if !errors.Is(err, ErrRequestCancelled) && !errors.Is(err, context.Canceled) {

98+ // Log the error (avoid logging cancellations as they're expected)

99+ logging.ErrorPersist(fmt.Sprintf("Generation error for session %s: %v", sessionID, err))

100+

101+ // You may want to create an error message in the chat

102+ bgCtx := context.Background()

103+ errorMsg := fmt.Sprintf("Sorry, an error occurred: %v", err)

104+ _, createErr := a.messages.Create(bgCtx, sessionID, message.CreateMessageParams{

105+ Role: message.System,

106+ Parts: []message.ContentPart{

107+ message.TextContent{

108+ Text: errorMsg,

109+ },

110+ },

111+ })

112+ if createErr != nil {

113+ logging.ErrorPersist(fmt.Sprintf("Failed to create error message: %v", createErr))

114+ }

115+ }

116+ }

117+ }()

118+

119+ return nil

120+}

121+

122+// IsSessionBusy checks if a session currently has an active request

123+func (a *agent) IsSessionBusy(sessionID string) bool {

124+ _, busy := a.activeRequests.Load(sessionID)

125+ return busy

126+} // handleTitleGeneration asynchronously generates a title for new sessions

32−func (c *agent) handleTitleGeneration(ctx context.Context, sessionID, content string) {

127+func (a *agent) handleTitleGeneration(ctx context.Context, sessionID, content string) {

33− response, err := c.titleGenerator.SendMessages(

128+ response, err := a.titleGenerator.SendMessages(

34 ctx,

129 ctx,

35 []message.Message{

130 []message.Message{

36 {

131 {

@@ -45,25 +140,30 @@ func (c *agent) handleTitleGeneration(ctx context.Context, sessionID, content st

45 nil,

140 nil,

46 )

141 )

47 if err != nil {

142 if err != nil {

143+ logging.ErrorPersist(fmt.Sprintf("Failed to generate title: %v", err))

48 return

144 return

49 }

145 }

146

51− session, err := c.Sessions.Get(ctx, sessionID)

147+ session, err := a.sessions.Get(ctx, sessionID)

52 if err != nil {

148 if err != nil {

149+ logging.ErrorPersist(fmt.Sprintf("Failed to get session: %v", err))

53 return

150 return

54 }

151 }

152+

55 if response.Content != "" {

153 if response.Content != "" {

56− session.Title = response.Content

154+ session.Title = strings.TrimSpace(response.Content)

57− session.Title = strings.TrimSpace(session.Title)

58 session.Title = strings.ReplaceAll(session.Title, "\n", " ")

155 session.Title = strings.ReplaceAll(session.Title, "\n", " ")

59− c.Sessions.Save(ctx, session)

156+ if _, err := a.sessions.Save(ctx, session); err != nil {

157+ logging.ErrorPersist(fmt.Sprintf("Failed to save session title: %v", err))

158+ }

60 }

159 }

61}

160}

161

162+// TrackUsage updates token usage statistics for the session

63−func (c *agent) TrackUsage(ctx context.Context, sessionID string, model models.Model, usage provider.TokenUsage) error {

163+func (a *agent) TrackUsage(ctx context.Context, sessionID string, model models.Model, usage provider.TokenUsage) error {

64− session, err := c.Sessions.Get(ctx, sessionID)

164+ session, err := a.sessions.Get(ctx, sessionID)

65 if err != nil {

165 if err != nil {

66− return err

166+ return fmt.Errorf("failed to get session: %w", err)

67 }

167 }

168

69 cost := model.CostPer1MInCached/1e6*float64(usage.CacheCreationTokens) +

169 cost := model.CostPer1MInCached/1e6*float64(usage.CacheCreationTokens) +

@@ -75,189 +175,241 @@ func (c *agent) TrackUsage(ctx context.Context, sessionID string, model models.M

75 session.CompletionTokens += usage.OutputTokens

175 session.CompletionTokens += usage.OutputTokens

76 session.PromptTokens += usage.InputTokens

176 session.PromptTokens += usage.InputTokens

177

78− _, err = c.Sessions.Save(ctx, session)

178+ _, err = a.sessions.Save(ctx, session)

79− return err

179+ if err != nil {

180+ return fmt.Errorf("failed to save session: %w", err)

181+ }

182+ return nil

80}

183}

184

185+// processEvent handles different types of events during generation

82−func (c *agent) processEvent(

186+func (a *agent) processEvent(

83 ctx context.Context,

187 ctx context.Context,

84 sessionID string,

188 sessionID string,

85 assistantMsg *message.Message,

189 assistantMsg *message.Message,

86 event provider.ProviderEvent,

190 event provider.ProviderEvent,

87) error {

191) error {

192+ select {

193+ case <-ctx.Done():

194+ return ctx.Err()

195+ default:

196+ // Continue processing

197+ }

198+

88 switch event.Type {

199 switch event.Type {

89 case provider.EventThinkingDelta:

200 case provider.EventThinkingDelta:

90 assistantMsg.AppendReasoningContent(event.Content)

201 assistantMsg.AppendReasoningContent(event.Content)

91− return c.Messages.Update(ctx, *assistantMsg)

202+ return a.messages.Update(ctx, *assistantMsg)

92 case provider.EventContentDelta:

203 case provider.EventContentDelta:

93 assistantMsg.AppendContent(event.Content)

204 assistantMsg.AppendContent(event.Content)

94− return c.Messages.Update(ctx, *assistantMsg)

205+ return a.messages.Update(ctx, *assistantMsg)

95 case provider.EventError:

206 case provider.EventError:

96 if errors.Is(event.Error, context.Canceled) {

207 if errors.Is(event.Error, context.Canceled) {

208+ logging.InfoPersist(fmt.Sprintf("Event processing canceled for session: %s", sessionID))

97− return nil

209+ return context.Canceled

98 }

210 }

99 logging.ErrorPersist(event.Error.Error())

211 logging.ErrorPersist(event.Error.Error())

100 return event.Error

212 return event.Error

101 case provider.EventWarning:

213 case provider.EventWarning:

102 logging.WarnPersist(event.Info)

214 logging.WarnPersist(event.Info)

103− return nil

104 case provider.EventInfo:

215 case provider.EventInfo:

105 logging.InfoPersist(event.Info)

216 logging.InfoPersist(event.Info)

106 case provider.EventComplete:

217 case provider.EventComplete:

107 assistantMsg.SetToolCalls(event.Response.ToolCalls)

218 assistantMsg.SetToolCalls(event.Response.ToolCalls)

108 assistantMsg.AddFinish(event.Response.FinishReason)

219 assistantMsg.AddFinish(event.Response.FinishReason)

109− err := c.Messages.Update(ctx, *assistantMsg)

220+ if err := a.messages.Update(ctx, *assistantMsg); err != nil {

110− if err != nil {

221+ return fmt.Errorf("failed to update message: %w", err)

111− return err

112 }

222 }

113− return c.TrackUsage(ctx, sessionID, c.model, event.Response.Usage)

223+ return a.TrackUsage(ctx, sessionID, a.model, event.Response.Usage)

114 }

224 }

115

225

116 return nil

226 return nil

117}

227}

118

228

229+// ExecuteTools runs all tool calls sequentially and returns the results

119−func (c *agent) ExecuteTools(ctx context.Context, toolCalls []message.ToolCall, tls []tools.BaseTool) ([]message.ToolResult, error) {

230+func (a *agent) ExecuteTools(ctx context.Context, toolCalls []message.ToolCall, tls []tools.BaseTool) ([]message.ToolResult, error) {

120− var wg sync.WaitGroup

121 toolResults := make([]message.ToolResult, len(toolCalls))

231 toolResults := make([]message.ToolResult, len(toolCalls))

122− mutex := &sync.Mutex{}

123− errChan := make(chan error, 1)

124

232

125 // Create a child context that can be canceled

233 // Create a child context that can be canceled

126 ctx, cancel := context.WithCancel(ctx)

234 ctx, cancel := context.WithCancel(ctx)

127 defer cancel()

235 defer cancel()

128

236

237+ // Check if already canceled before starting any execution

238+ if ctx.Err() != nil {

239+ // Mark all tools as canceled

129− for i, tc := range toolCalls {

240+ for i, toolCall := range toolCalls {

130− wg.Add(1)

241+ toolResults[i] = message.ToolResult{

131− go func(index int, toolCall message.ToolCall) {

242+ ToolCallID: toolCall.ID,

132− defer wg.Done()

243+ Content: "Tool execution canceled by user",

244+ IsError: true,

245+ }

246+ }

247+ return toolResults, ctx.Err()

248+ }

133

249

134− // Check if context is already canceled

250+ for i, toolCall := range toolCalls {

251+ // Check for cancellation before executing each tool

135− select {

252+ select {

136− case <-ctx.Done():

253+ case <-ctx.Done():

137− mutex.Lock()

254+ // Mark this and all remaining tools as canceled

255+ for j := i; j < len(toolCalls); j++ {

138− toolResults[index] = message.ToolResult{

256+ toolResults[j] = message.ToolResult{

139− ToolCallID: toolCall.ID,

257+ ToolCallID: toolCalls[j].ID,

140− Content: "Tool execution canceled",

258+ Content: "Tool execution canceled by user",

141 IsError: true,

259 IsError: true,

142 }

260 }

143− mutex.Unlock()

144−

145− // Send cancellation error to error channel if it's empty

146− select {

147− case errChan <- ctx.Err():

148− default:

149− }

150− return

151− default:

152 }

261 }

262+ return toolResults, ctx.Err()

263+ default:

264+ // Continue processing

265+ }

153

266

154− response := ""

267+ response := ""

155− isError := false

268+ isError := false

156− found := false

269+ found := false

157−

270+

271+ // Find and execute the appropriate tool

158− for _, tool := range tls {

272+ for _, tool := range tls {

159− if tool.Info().Name == toolCall.Name {

273+ if tool.Info().Name == toolCall.Name {

160− found = true

274+ found = true

161− toolResult, toolErr := tool.Run(ctx, tools.ToolCall{

275+ toolResult, toolErr := tool.Run(ctx, tools.ToolCall{

162− ID: toolCall.ID,

276+ ID: toolCall.ID,

163− Name: toolCall.Name,

277+ Name: toolCall.Name,

164− Input: toolCall.Input,

278+ Input: toolCall.Input,

165− })

279+ })

166−

280+

167− if toolErr != nil {

281+ if toolErr != nil {

168− if errors.Is(toolErr, context.Canceled) {

282+ if errors.Is(toolErr, context.Canceled) {

169− response = "Tool execution canceled"

283+ response = "Tool execution canceled by user"

170−

171− // Send cancellation error to error channel if it's empty

172− select {

173− case errChan <- ctx.Err():

174− default:

175− }

176− } else {

177− response = fmt.Sprintf("error running tool: %s", toolErr)

178− }

179− isError = true

180 } else {

284 } else {

181− response = toolResult.Content

285+ response = fmt.Sprintf("Error running tool: %s", toolErr)

182− isError = toolResult.IsError

183 }

286 }

184− break

287+ isError = true

288+ } else {

289+ response = toolResult.Content

290+ isError = toolResult.IsError

185 }

291 }

292+ break

186 }

293 }

294+ }

187

295

188− if !found {

296+ if !found {

189− response = fmt.Sprintf("tool not found: %s", toolCall.Name)

297+ response = fmt.Sprintf("Tool not found: %s", toolCall.Name)

190− isError = true

298+ isError = true

191− }

192−

193− mutex.Lock()

194− defer mutex.Unlock()

195−

196− toolResults[index] = message.ToolResult{

197− ToolCallID: toolCall.ID,

198− Content: response,

199− IsError: isError,

200− }

201− }(i, tc)

202− }

203−

204− // Wait for all goroutines to finish or context to be canceled

205− done := make(chan struct{})

206− go func() {

207− wg.Wait()

208− close(done)

209− }()

299+ }

210

300

211− select {

301+ toolResults[i] = message.ToolResult{

212− case <-done:

302+ ToolCallID: toolCall.ID,

213− // All tools completed successfully

303+ Content: response,

214− case err := <-errChan:

304+ IsError: isError,

215− // One of the tools encountered a cancellation

305+ }

216− return toolResults, err

217− case <-ctx.Done():

218− // Context was canceled externally

219− return toolResults, ctx.Err()

220 }

306 }

221

307

222 return toolResults, nil

308 return toolResults, nil

223}

309}

224

310

311+// handleToolExecution processes tool calls and creates tool result messages

225−func (c *agent) handleToolExecution(

312+func (a *agent) handleToolExecution(

226 ctx context.Context,

313 ctx context.Context,

227 assistantMsg message.Message,

314 assistantMsg message.Message,

228) (*message.Message, error) {

315) (*message.Message, error) {

316+ select {

317+ case <-ctx.Done():

318+ // If cancelled, create tool results that indicate cancellation

319+ if len(assistantMsg.ToolCalls()) > 0 {

320+ toolResults := make([]message.ToolResult, 0, len(assistantMsg.ToolCalls()))

321+ for _, tc := range assistantMsg.ToolCalls() {

322+ toolResults = append(toolResults, message.ToolResult{

323+ ToolCallID: tc.ID,

324+ Content: "Tool execution canceled by user",

325+ IsError: true,

326+ })

327+ }

328+

329+ // Use background context to ensure the message is created even if original context is cancelled

330+ bgCtx := context.Background()

331+ parts := make([]message.ContentPart, 0)

332+ for _, toolResult := range toolResults {

333+ parts = append(parts, toolResult)

334+ }

335+ msg, err := a.messages.Create(bgCtx, assistantMsg.SessionID, message.CreateMessageParams{

336+ Role: message.Tool,

337+ Parts: parts,

338+ })

339+ if err != nil {

340+ return nil, fmt.Errorf("failed to create cancelled tool message: %w", err)

341+ }

342+ return &msg, ctx.Err()

343+ }

344+ return nil, ctx.Err()

345+ default:

346+ // Continue processing

347+ }

348+

229 if len(assistantMsg.ToolCalls()) == 0 {

349 if len(assistantMsg.ToolCalls()) == 0 {

230 return nil, nil

350 return nil, nil

231 }

351 }

232

352

233− toolResults, err := c.ExecuteTools(ctx, assistantMsg.ToolCalls(), c.tools)

353+ toolResults, err := a.ExecuteTools(ctx, assistantMsg.ToolCalls(), a.tools)

234 if err != nil {

354 if err != nil {

355+ // If error is from cancellation, still return the partial results we have

356+ if errors.Is(err, context.Canceled) {

357+ // Use background context to ensure the message is created even if original context is cancelled

358+ bgCtx := context.Background()

359+ parts := make([]message.ContentPart, 0)

360+ for _, toolResult := range toolResults {

361+ parts = append(parts, toolResult)

362+ }

363+

364+ msg, createErr := a.messages.Create(bgCtx, assistantMsg.SessionID, message.CreateMessageParams{

365+ Role: message.Tool,

366+ Parts: parts,

367+ })

368+ if createErr != nil {

369+ logging.ErrorPersist(fmt.Sprintf("Failed to create tool message after cancellation: %v", createErr))

370+ return nil, err

371+ }

372+ return &msg, err

373+ }

235 return nil, err

374 return nil, err

236 }

375 }

376+

237− parts := make([]message.ContentPart, 0)

377+ parts := make([]message.ContentPart, 0, len(toolResults))

238 for _, toolResult := range toolResults {

378 for _, toolResult := range toolResults {

239 parts = append(parts, toolResult)

379 parts = append(parts, toolResult)

240 }

380 }

381+

241− msg, err := c.Messages.Create(ctx, assistantMsg.SessionID, message.CreateMessageParams{

382+ msg, err := a.messages.Create(ctx, assistantMsg.SessionID, message.CreateMessageParams{

242 Role: message.Tool,

383 Role: message.Tool,

243 Parts: parts,

384 Parts: parts,

244 })

385 })

386+ if err != nil {

387+ return nil, fmt.Errorf("failed to create tool message: %w", err)

388+ }

245

389

246− return &msg, err

390+ return &msg, nil

247}

391}

248

392

393+// generate handles the main generation workflow

249−func (c *agent) generate(ctx context.Context, sessionID string, content string) error {

394+func (a *agent) generate(ctx context.Context, sessionID string, content string) error {

250 ctx = context.WithValue(ctx, tools.SessionIDContextKey, sessionID)

395 ctx = context.WithValue(ctx, tools.SessionIDContextKey, sessionID)

396+

397+ // Handle context cancellation at any point

398+ if err := ctx.Err(); err != nil {

399+ return ErrRequestCancelled

400+ }

401+

251− messages, err := c.Messages.List(ctx, sessionID)

402+ messages, err := a.messages.List(ctx, sessionID)

252 if err != nil {

403 if err != nil {

253− return err

404+ return fmt.Errorf("failed to list messages: %w", err)

254 }

405 }

255

406

256 if len(messages) == 0 {

407 if len(messages) == 0 {

408+ titleCtx := context.Background()

257− go c.handleTitleGeneration(ctx, sessionID, content)

409+ go a.handleTitleGeneration(titleCtx, sessionID, content)

258 }

410 }

259

411

260− userMsg, err := c.Messages.Create(ctx, sessionID, message.CreateMessageParams{

412+ userMsg, err := a.messages.Create(ctx, sessionID, message.CreateMessageParams{

261 Role: message.User,

413 Role: message.User,

262 Parts: []message.ContentPart{

414 Parts: []message.ContentPart{

263 message.TextContent{

415 message.TextContent{

@@ -266,133 +418,125 @@ func (c *agent) generate(ctx context.Context, sessionID string, content string)

266 },

418 },

267 })

419 })

268 if err != nil {

420 if err != nil {

269− return err

421+ return fmt.Errorf("failed to create user message: %w", err)

270 }

422 }

271

423

272 messages = append(messages, userMsg)

424 messages = append(messages, userMsg)

425+

273 for {

426 for {

427+ // Check for cancellation before each iteration

274 select {

428 select {

275 case <-ctx.Done():

429 case <-ctx.Done():

276− assistantMsg, err := c.Messages.Create(ctx, sessionID, message.CreateMessageParams{

430+ return ErrRequestCancelled

277− Role: message.Assistant,

278− Parts: []message.ContentPart{},

279− })

280− if err != nil {

281− return err

282− }

283− assistantMsg.AddFinish("canceled")

284− c.Messages.Update(ctx, assistantMsg)

285− return context.Canceled

286 default:

431 default:

287 // Continue processing

432 // Continue processing

288 }

433 }

289

434

290− eventChan, err := c.agent.StreamResponse(ctx, messages, c.tools)

435+ eventChan, err := a.agent.StreamResponse(ctx, messages, a.tools)

291 if err != nil {

436 if err != nil {

292 if errors.Is(err, context.Canceled) {

437 if errors.Is(err, context.Canceled) {

293− assistantMsg, err := c.Messages.Create(ctx, sessionID, message.CreateMessageParams{

438+ return ErrRequestCancelled

294− Role: message.Assistant,

295− Parts: []message.ContentPart{},

296− })

297− if err != nil {

298− return err

299− }

300− assistantMsg.AddFinish("canceled")

301− c.Messages.Update(ctx, assistantMsg)

302− return context.Canceled

303 }

439 }

304− return err

440+ return fmt.Errorf("failed to stream response: %w", err)

305 }

441 }

306

442

307− assistantMsg, err := c.Messages.Create(ctx, sessionID, message.CreateMessageParams{

443+ assistantMsg, err := a.messages.Create(ctx, sessionID, message.CreateMessageParams{

308 Role: message.Assistant,

444 Role: message.Assistant,

309 Parts: []message.ContentPart{},

445 Parts: []message.ContentPart{},

310− Model: c.model.ID,

446+ Model: a.model.ID,

311 })

447 })

312 if err != nil {

448 if err != nil {

313− return err

449+ return fmt.Errorf("failed to create assistant message: %w", err)

314 }

450 }

315

451

316 ctx = context.WithValue(ctx, tools.MessageIDContextKey, assistantMsg.ID)

452 ctx = context.WithValue(ctx, tools.MessageIDContextKey, assistantMsg.ID)

453+

454+ // Process events from the LLM provider

317 for event := range eventChan {

455 for event := range eventChan {

318− err = c.processEvent(ctx, sessionID, &assistantMsg, event)

456+ if err := a.processEvent(ctx, sessionID, &assistantMsg, event); err != nil {

319− if err != nil {

320 if errors.Is(err, context.Canceled) {

457 if errors.Is(err, context.Canceled) {

458+ // Mark as canceled but don't create separate message

321 assistantMsg.AddFinish("canceled")

459 assistantMsg.AddFinish("canceled")

322− c.Messages.Update(ctx, assistantMsg)

460+ _ = a.messages.Update(context.Background(), assistantMsg)

323− return context.Canceled

461+ return ErrRequestCancelled

324 }

462 }

325 assistantMsg.AddFinish("error:" + err.Error())

463 assistantMsg.AddFinish("error:" + err.Error())

326− c.Messages.Update(ctx, assistantMsg)

464+ _ = a.messages.Update(ctx, assistantMsg)

327− return err

465+ return fmt.Errorf("event processing error: %w", err)

328 }

466 }

329

467

468+ // Check for cancellation during event processing

330 select {

469 select {

331 case <-ctx.Done():

470 case <-ctx.Done():

471+ // Mark as canceled

332 assistantMsg.AddFinish("canceled")

472 assistantMsg.AddFinish("canceled")

333− c.Messages.Update(ctx, assistantMsg)

473+ _ = a.messages.Update(context.Background(), assistantMsg)

334− return context.Canceled

474+ return ErrRequestCancelled

335 default:

475 default:

336 }

476 }

337 }

477 }

338

478

339− // Check for context cancellation before tool execution

479+ // Check for cancellation before tool execution

340 select {

480 select {

341 case <-ctx.Done():

481 case <-ctx.Done():

342− assistantMsg.AddFinish("canceled")

482+ assistantMsg.AddFinish("canceled_by_user")

343− c.Messages.Update(ctx, assistantMsg)

483+ _ = a.messages.Update(context.Background(), assistantMsg)

344− return context.Canceled

484+ return ErrRequestCancelled

345 default:

485 default:

346− // Continue processing

347 }

486 }

348

487

488+ // Execute any tool calls

349− msg, err := c.handleToolExecution(ctx, assistantMsg)

489+ toolMsg, err := a.handleToolExecution(ctx, assistantMsg)

350 if err != nil {

490 if err != nil {

351 if errors.Is(err, context.Canceled) {

491 if errors.Is(err, context.Canceled) {

352− assistantMsg.AddFinish("canceled")

492+ assistantMsg.AddFinish("canceled_by_user")

353− c.Messages.Update(ctx, assistantMsg)

493+ _ = a.messages.Update(context.Background(), assistantMsg)

354− return context.Canceled

494+ return ErrRequestCancelled

355 }

495 }

356− return err

496+ return fmt.Errorf("tool execution error: %w", err)

357 }

497 }

358

498

359− c.Messages.Update(ctx, assistantMsg)

499+ if err := a.messages.Update(ctx, assistantMsg); err != nil {

500+ return fmt.Errorf("failed to update assistant message: %w", err)

501+ }

360

502

503+ // If no tool calls, we're done

361 if len(assistantMsg.ToolCalls()) == 0 {

504 if len(assistantMsg.ToolCalls()) == 0 {

362 break

505 break

363 }

506 }

364

507

508+ // Add messages for next iteration

365 messages = append(messages, assistantMsg)

509 messages = append(messages, assistantMsg)

366− if msg != nil {

510+ if toolMsg != nil {

367− messages = append(messages, *msg)

511+ messages = append(messages, *toolMsg)

368 }

512 }

369

513

370− // Check for context cancellation after tool execution

514+ // Check for cancellation after tool execution

371 select {

515 select {

372 case <-ctx.Done():

516 case <-ctx.Done():

373− assistantMsg.AddFinish("canceled")

517+ return ErrRequestCancelled

374− c.Messages.Update(ctx, assistantMsg)

375− return context.Canceled

376 default:

518 default:

377− // Continue processing

378 }

519 }

379 }

520 }

521+

380 return nil

522 return nil

381}

523}

382

524

525+// getAgentProviders initializes the LLM providers based on the chosen model

383func getAgentProviders(ctx context.Context, model models.Model) (provider.Provider, provider.Provider, error) {

526func getAgentProviders(ctx context.Context, model models.Model) (provider.Provider, provider.Provider, error) {

384 maxTokens := config.Get().Model.CoderMaxTokens

527 maxTokens := config.Get().Model.CoderMaxTokens

385

528

386 providerConfig, ok := config.Get().Providers[model.Provider]

529 providerConfig, ok := config.Get().Providers[model.Provider]

387 if !ok || providerConfig.Disabled {

530 if !ok || providerConfig.Disabled {

388− return nil, nil, errors.New("provider is not enabled")

531+ return nil, nil, ErrProviderNotEnabled

389 }

532 }

533+

390 var agentProvider provider.Provider

534 var agentProvider provider.Provider

391 var titleGenerator provider.Provider

535 var titleGenerator provider.Provider

536+ var err error

392

537

393 switch model.Provider {

538 switch model.Provider {

394 case models.ProviderOpenAI:

539 case models.ProviderOpenAI:

395− var err error

396 agentProvider, err = provider.NewOpenAIProvider(

540 agentProvider, err = provider.NewOpenAIProvider(

397 provider.WithOpenAISystemMessage(

541 provider.WithOpenAISystemMessage(

398 prompt.CoderOpenAISystemPrompt(),

542 prompt.CoderOpenAISystemPrompt(),

@@ -402,8 +546,9 @@ func getAgentProviders(ctx context.Context, model models.Model) (provider.Provid

402 provider.WithOpenAIKey(providerConfig.APIKey),

546 provider.WithOpenAIKey(providerConfig.APIKey),

403 )

547 )

404 if err != nil {

548 if err != nil {

405− return nil, nil, err

549+ return nil, nil, fmt.Errorf("failed to create OpenAI agent provider: %w", err)

406 }

550 }

551+

407 titleGenerator, err = provider.NewOpenAIProvider(

552 titleGenerator, err = provider.NewOpenAIProvider(

408 provider.WithOpenAISystemMessage(

553 provider.WithOpenAISystemMessage(

409 prompt.TitlePrompt(),

554 prompt.TitlePrompt(),

@@ -413,10 +558,10 @@ func getAgentProviders(ctx context.Context, model models.Model) (provider.Provid

413 provider.WithOpenAIKey(providerConfig.APIKey),

558 provider.WithOpenAIKey(providerConfig.APIKey),

414 )

559 )

415 if err != nil {

560 if err != nil {

416− return nil, nil, err

561+ return nil, nil, fmt.Errorf("failed to create OpenAI title generator: %w", err)

417 }

562 }

563+

418 case models.ProviderAnthropic:

564 case models.ProviderAnthropic:

419− var err error

420 agentProvider, err = provider.NewAnthropicProvider(

565 agentProvider, err = provider.NewAnthropicProvider(

421 provider.WithAnthropicSystemMessage(

566 provider.WithAnthropicSystemMessage(

422 prompt.CoderAnthropicSystemPrompt(),

567 prompt.CoderAnthropicSystemPrompt(),

@@ -426,8 +571,9 @@ func getAgentProviders(ctx context.Context, model models.Model) (provider.Provid

426 provider.WithAnthropicModel(model),

571 provider.WithAnthropicModel(model),

427 )

572 )

428 if err != nil {

573 if err != nil {

429− return nil, nil, err

574+ return nil, nil, fmt.Errorf("failed to create Anthropic agent provider: %w", err)

430 }

575 }

576+

431 titleGenerator, err = provider.NewAnthropicProvider(

577 titleGenerator, err = provider.NewAnthropicProvider(

432 provider.WithAnthropicSystemMessage(

578 provider.WithAnthropicSystemMessage(

433 prompt.TitlePrompt(),

579 prompt.TitlePrompt(),

@@ -437,11 +583,10 @@ func getAgentProviders(ctx context.Context, model models.Model) (provider.Provid

437 provider.WithAnthropicModel(model),

583 provider.WithAnthropicModel(model),

438 )

584 )

439 if err != nil {

585 if err != nil {

440− return nil, nil, err

586+ return nil, nil, fmt.Errorf("failed to create Anthropic title generator: %w", err)

441 }

587 }

442

588

443 case models.ProviderGemini:

589 case models.ProviderGemini:

444− var err error

445 agentProvider, err = provider.NewGeminiProvider(

590 agentProvider, err = provider.NewGeminiProvider(

446 ctx,

591 ctx,

447 provider.WithGeminiSystemMessage(

592 provider.WithGeminiSystemMessage(

@@ -452,8 +597,9 @@ func getAgentProviders(ctx context.Context, model models.Model) (provider.Provid

452 provider.WithGeminiModel(model),

597 provider.WithGeminiModel(model),

453 )

598 )

454 if err != nil {

599 if err != nil {

455− return nil, nil, err

600+ return nil, nil, fmt.Errorf("failed to create Gemini agent provider: %w", err)

456 }

601 }

602+

457 titleGenerator, err = provider.NewGeminiProvider(

603 titleGenerator, err = provider.NewGeminiProvider(

458 ctx,

604 ctx,

459 provider.WithGeminiSystemMessage(

605 provider.WithGeminiSystemMessage(

@@ -464,10 +610,10 @@ func getAgentProviders(ctx context.Context, model models.Model) (provider.Provid

464 provider.WithGeminiModel(model),

610 provider.WithGeminiModel(model),

465 )

611 )

466 if err != nil {

612 if err != nil {

467− return nil, nil, err

613+ return nil, nil, fmt.Errorf("failed to create Gemini title generator: %w", err)

468 }

614 }

615+

469 case models.ProviderGROQ:

616 case models.ProviderGROQ:

470− var err error

471 agentProvider, err = provider.NewOpenAIProvider(

617 agentProvider, err = provider.NewOpenAIProvider(

472 provider.WithOpenAISystemMessage(

618 provider.WithOpenAISystemMessage(

473 prompt.CoderAnthropicSystemPrompt(),

619 prompt.CoderAnthropicSystemPrompt(),

@@ -478,8 +624,9 @@ func getAgentProviders(ctx context.Context, model models.Model) (provider.Provid

478 provider.WithOpenAIBaseURL("https://api.groq.com/openai/v1"),

624 provider.WithOpenAIBaseURL("https://api.groq.com/openai/v1"),

479 )

625 )

480 if err != nil {

626 if err != nil {

481− return nil, nil, err

627+ return nil, nil, fmt.Errorf("failed to create GROQ agent provider: %w", err)

482 }

628 }

629+

483 titleGenerator, err = provider.NewOpenAIProvider(

630 titleGenerator, err = provider.NewOpenAIProvider(

484 provider.WithOpenAISystemMessage(

631 provider.WithOpenAISystemMessage(

485 prompt.TitlePrompt(),

632 prompt.TitlePrompt(),

@@ -490,11 +637,10 @@ func getAgentProviders(ctx context.Context, model models.Model) (provider.Provid

490 provider.WithOpenAIBaseURL("https://api.groq.com/openai/v1"),

637 provider.WithOpenAIBaseURL("https://api.groq.com/openai/v1"),

491 )

638 )

492 if err != nil {

639 if err != nil {

493− return nil, nil, err

640+ return nil, nil, fmt.Errorf("failed to create GROQ title generator: %w", err)

494 }

641 }

495

642

496 case models.ProviderBedrock:

643 case models.ProviderBedrock:

497− var err error

498 agentProvider, err = provider.NewBedrockProvider(

644 agentProvider, err = provider.NewBedrockProvider(

499 provider.WithBedrockSystemMessage(

645 provider.WithBedrockSystemMessage(

500 prompt.CoderAnthropicSystemPrompt(),

646 prompt.CoderAnthropicSystemPrompt(),

@@ -503,19 +649,21 @@ func getAgentProviders(ctx context.Context, model models.Model) (provider.Provid

503 provider.WithBedrockModel(model),

649 provider.WithBedrockModel(model),

504 )

650 )

505 if err != nil {

651 if err != nil {

506− return nil, nil, err

652+ return nil, nil, fmt.Errorf("failed to create Bedrock agent provider: %w", err)

507 }

653 }

654+

508 titleGenerator, err = provider.NewBedrockProvider(

655 titleGenerator, err = provider.NewBedrockProvider(

509 provider.WithBedrockSystemMessage(

656 provider.WithBedrockSystemMessage(

510 prompt.TitlePrompt(),

657 prompt.TitlePrompt(),

511 ),

658 ),

512− provider.WithBedrockMaxTokens(maxTokens),

659+ provider.WithBedrockMaxTokens(80),

513 provider.WithBedrockModel(model),

660 provider.WithBedrockModel(model),

514 )

661 )

515 if err != nil {

662 if err != nil {

516− return nil, nil, err

663+ return nil, nil, fmt.Errorf("failed to create Bedrock title generator: %w", err)

517 }

664 }

518−

665+ default:

666+ return nil, nil, fmt.Errorf("unsupported provider: %s", model.Provider)

519 }

667 }

520

668

521 return agentProvider, titleGenerator, nil

669 return agentProvider, titleGenerator, nil

internal/llm/agent/coder.go+36−47

@@ -4,71 +4,60 @@ import (

4 "context"

5 "errors"

7− "github.com/kujtimiihoxha/termai/internal/app"

8 "github.com/kujtimiihoxha/termai/internal/config"

7 "github.com/kujtimiihoxha/termai/internal/config"

9 "github.com/kujtimiihoxha/termai/internal/llm/models"

8 "github.com/kujtimiihoxha/termai/internal/llm/models"

10 "github.com/kujtimiihoxha/termai/internal/llm/tools"

9 "github.com/kujtimiihoxha/termai/internal/llm/tools"

10+ "github.com/kujtimiihoxha/termai/internal/lsp"

11+ "github.com/kujtimiihoxha/termai/internal/message"

12+ "github.com/kujtimiihoxha/termai/internal/permission"

13+ "github.com/kujtimiihoxha/termai/internal/session"

11)

14)

13type coderAgent struct {

16type coderAgent struct {

14− *agent

17+ Service

15}

18}

17−func (c *coderAgent) setAgentTool(sessionID string) {

20+func NewCoderAgent(

18− inx := -1

21+ permissions permission.Service,

19− for i, tool := range c.tools {

22+ sessions session.Service,

20− if tool.Info().Name == AgentToolName {

23+ messages message.Service,

21− inx = i

24+ lspClients map[string]*lsp.Client,

22− break

25+) (Service, error) {

23− }

24− }

25− if inx == -1 {

26− c.tools = append(c.tools, NewAgentTool(sessionID, c.App))

27− } else {

28− c.tools[inx] = NewAgentTool(sessionID, c.App)

29− }

30−}

31−

32−func (c *coderAgent) Generate(ctx context.Context, sessionID string, content string) error {

33− c.setAgentTool(sessionID)

34− return c.generate(ctx, sessionID, content)

35−}

36−

37−func NewCoderAgent(app *app.App) (Agent, error) {

38 model, ok := models.SupportedModels[config.Get().Model.Coder]

26 model, ok := models.SupportedModels[config.Get().Model.Coder]

39 if !ok {

27 if !ok {

40 return nil, errors.New("model not supported")

28 return nil, errors.New("model not supported")

41 }

29 }

43 ctx := context.Background()

31 ctx := context.Background()

44− agentProvider, titleGenerator, err := getAgentProviders(ctx, model)

32+ otherTools := GetMcpTools(ctx, permissions)

33+ if len(lspClients) > 0 {

34+ otherTools = append(otherTools, tools.NewDiagnosticsTool(lspClients))

35+ }

36+ agent, err := NewAgent(

37+ ctx,

38+ sessions,

39+ messages,

40+ model,

41+ append(

42+ []tools.BaseTool{

43+ tools.NewBashTool(permissions),

44+ tools.NewEditTool(lspClients, permissions),

45+ tools.NewFetchTool(permissions),

46+ tools.NewGlobTool(),

47+ tools.NewGrepTool(),

48+ tools.NewLsTool(),

49+ tools.NewSourcegraphTool(),

50+ tools.NewViewTool(lspClients),

51+ tools.NewWriteTool(lspClients, permissions),

52+ NewAgentTool(sessions, messages),

53+ }, otherTools...,

54+ ),

55+ )

45 if err != nil {

56 if err != nil {

46 return nil, err

57 return nil, err

47 }

58 }

49− otherTools := GetMcpTools(ctx, app.Permissions)

50− if len(app.LSPClients) > 0 {

51− otherTools = append(otherTools, tools.NewDiagnosticsTool(app.LSPClients))

52− }

53 return &coderAgent{

60 return &coderAgent{

54− agent: &agent{

61+ agent,

55− App: app,

56− tools: append(

57− []tools.BaseTool{

58− tools.NewBashTool(app.Permissions),

59− tools.NewEditTool(app.LSPClients, app.Permissions),

60− tools.NewFetchTool(app.Permissions),

61− tools.NewGlobTool(),

62− tools.NewGrepTool(),

63− tools.NewLsTool(),

64− tools.NewSourcegraphTool(),

65− tools.NewViewTool(app.LSPClients),

66− tools.NewWriteTool(app.LSPClients, app.Permissions),

67− }, otherTools...,

68− ),

69− model: model,

70− agent: agentProvider,

71− titleGenerator: titleGenerator,

72− },

73 }, nil

62 }, nil

74}

63}

internal/llm/agent/task.go+3−4

@@ -4,10 +4,10 @@ import (

4 "context"

5 "errors"

7− "github.com/kujtimiihoxha/termai/internal/app"

8 "github.com/kujtimiihoxha/termai/internal/config"

7 "github.com/kujtimiihoxha/termai/internal/config"

9 "github.com/kujtimiihoxha/termai/internal/llm/models"

8 "github.com/kujtimiihoxha/termai/internal/llm/models"

10 "github.com/kujtimiihoxha/termai/internal/llm/tools"

9 "github.com/kujtimiihoxha/termai/internal/llm/tools"

10+ "github.com/kujtimiihoxha/termai/internal/lsp"

11)

13type taskAgent struct {

@@ -18,7 +18,7 @@ func (c *taskAgent) Generate(ctx context.Context, sessionID string, content stri

18 return c.generate(ctx, sessionID, content)

19}

21−func NewTaskAgent(app *app.App) (Agent, error) {

21+func NewTaskAgent(lspClients map[string]*lsp.Client) (Service, error) {

22 model, ok := models.SupportedModels[config.Get().Model.Coder]

23 if !ok {

24 return nil, errors.New("model not supported")

@@ -31,13 +31,12 @@ func NewTaskAgent(app *app.App) (Agent, error) {

31 }

32 return &taskAgent{

33 agent: &agent{

34− App: app,

35 tools: []tools.BaseTool{

34 tools: []tools.BaseTool{

36 tools.NewGlobTool(),

35 tools.NewGlobTool(),

37 tools.NewGrepTool(),

36 tools.NewGrepTool(),

38 tools.NewLsTool(),

37 tools.NewLsTool(),

39 tools.NewSourcegraphTool(),

38 tools.NewSourcegraphTool(),

40− tools.NewViewTool(app.LSPClients),

39+ tools.NewViewTool(lspClients),

41 },

40 },

42 model: model,

41 model: model,

43 agent: agentProvider,

42 agent: agentProvider,

internal/llm/provider/provider.go+3−1

@@ -57,7 +57,9 @@ func cleanupMessages(messages []message.Message) []message.Message {

57 // First pass: filter out canceled messages

58 var cleanedMessages []message.Message

59 for _, msg := range messages {

60− if msg.FinishReason() != "canceled" {

60+ if msg.FinishReason() != "canceled" || len(msg.ToolCalls()) > 0 {

61+ // if there are toolCalls this means we want to return it to the LLM telling it that those tools have been

62+ // cancelled

61 cleanedMessages = append(cleanedMessages, msg)

63 cleanedMessages = append(cleanedMessages, msg)

62 }

64 }

63 }

65 }

internal/llm/tools/edit.go+3−4

@@ -190,7 +190,7 @@ func (e *editTool) createNewFile(ctx context.Context, filePath, content string)

190 return er, fmt.Errorf("failed to create parent directories: %w", err)

191 }

192

193− sessionID, messageID := getContextValues(ctx)

193+ sessionID, messageID := GetContextValues(ctx)

194 if sessionID == "" || messageID == "" {

195 return er, fmt.Errorf("session ID and message ID are required for creating a new file")

196 }

@@ -277,7 +277,7 @@ func (e *editTool) deleteContent(ctx context.Context, filePath, oldString string

277

278 newContent := oldContent[:index] + oldContent[index+len(oldString):]

279

280− sessionID, messageID := getContextValues(ctx)

280+ sessionID, messageID := GetContextValues(ctx)

281

282 if sessionID == "" || messageID == "" {

283 return er, fmt.Errorf("session ID and message ID are required for creating a new file")

@@ -365,7 +365,7 @@ func (e *editTool) replaceContent(ctx context.Context, filePath, oldString, newS

365

366 newContent := oldContent[:index] + newString + oldContent[index+len(oldString):]

367

368− sessionID, messageID := getContextValues(ctx)

368+ sessionID, messageID := GetContextValues(ctx)

369

370 if sessionID == "" || messageID == "" {

371 return er, fmt.Errorf("session ID and message ID are required for creating a new file")

@@ -409,4 +409,3 @@ func (e *editTool) replaceContent(ctx context.Context, filePath, oldString, newS

409

410 return er, nil

411}

412−

internal/llm/tools/tools.go+1−1

@@ -66,7 +66,7 @@ type BaseTool interface {

66 Run(ctx context.Context, params ToolCall) (ToolResponse, error)

67}

69−func getContextValues(ctx context.Context) (string, string) {

69+func GetContextValues(ctx context.Context) (string, string) {

70 sessionID := ctx.Value(SessionIDContextKey)

71 messageID := ctx.Value(MessageIDContextKey)

72 if sessionID == nil {

internal/llm/tools/write.go+1−1

@@ -144,7 +144,7 @@ func (w *writeTool) Run(ctx context.Context, call ToolCall) (ToolResponse, error

144 }

145 }

146

147− sessionID, messageID := getContextValues(ctx)

147+ sessionID, messageID := GetContextValues(ctx)

148 if sessionID == "" || messageID == "" {

149 return NewTextErrorResponse("session ID or message ID is missing"), nil

150 }

internal/tui/components/repl/editor.go+1−7

@@ -7,7 +7,6 @@ import (

7 tea "github.com/charmbracelet/bubbletea"

8 "github.com/charmbracelet/lipgloss"

9 "github.com/kujtimiihoxha/termai/internal/app"

10− "github.com/kujtimiihoxha/termai/internal/llm/agent"

11 "github.com/kujtimiihoxha/termai/internal/tui/layout"

10 "github.com/kujtimiihoxha/termai/internal/tui/layout"

12 "github.com/kujtimiihoxha/termai/internal/tui/styles"

11 "github.com/kujtimiihoxha/termai/internal/tui/styles"

13 "github.com/kujtimiihoxha/termai/internal/tui/util"

12 "github.com/kujtimiihoxha/termai/internal/tui/util"

@@ -168,11 +167,6 @@ func (m *editorCmp) Send() tea.Cmd {

168 return util.ReportWarn("Assistant is still working on the previous message")

167 return util.ReportWarn("Assistant is still working on the previous message")

169 }

168 }

170

169

171− a, err := agent.NewCoderAgent(m.app)

172− if err != nil {

173− return util.ReportError(err)

174− }

175−

176 content := strings.Join(m.editor.GetBuffer().Lines(), "\n")

170 content := strings.Join(m.editor.GetBuffer().Lines(), "\n")

177 if len(content) == 0 {

171 if len(content) == 0 {

178 return util.ReportWarn("Message is empty")

172 return util.ReportWarn("Message is empty")

@@ -181,7 +175,7 @@ func (m *editorCmp) Send() tea.Cmd {

181 m.cancelMessage = cancel

175 m.cancelMessage = cancel

182 go func() {

176 go func() {

183 defer cancel()

177 defer cancel()

184− a.Generate(ctx, m.sessionID, content)

178+ m.app.CoderAgent.Generate(ctx, m.sessionID, content)

185 m.cancelMessage = nil

179 m.cancelMessage = nil

186 }()

180 }()

187

181

internal/tui/page/chat.go+6−9

@@ -6,7 +6,6 @@ import (

6 "github.com/charmbracelet/bubbles/key"

7 tea "github.com/charmbracelet/bubbletea"

8 "github.com/kujtimiihoxha/termai/internal/app"

9− "github.com/kujtimiihoxha/termai/internal/llm/agent"

10 "github.com/kujtimiihoxha/termai/internal/session"

9 "github.com/kujtimiihoxha/termai/internal/session"

11 "github.com/kujtimiihoxha/termai/internal/tui/components/chat"

10 "github.com/kujtimiihoxha/termai/internal/tui/components/chat"

12 "github.com/kujtimiihoxha/termai/internal/tui/layout"

11 "github.com/kujtimiihoxha/termai/internal/tui/layout"

@@ -23,6 +22,7 @@ type chatPage struct {

24type ChatKeyMap struct {

23type ChatKeyMap struct {

25 NewSession key.Binding

24 NewSession key.Binding

25+ Cancel key.Binding

26}

28var keyMap = ChatKeyMap{

@@ -30,6 +30,10 @@ var keyMap = ChatKeyMap{

30 key.WithKeys("ctrl+n"),

31 key.WithHelp("ctrl+n", "new session"),

32 ),

33+ Cancel: key.NewBinding(

34+ key.WithKeys("ctrl+x"),

35+ key.WithHelp("ctrl+x", "cancel"),

36+ ),

33}

37}

35func (p *chatPage) Init() tea.Cmd {

39func (p *chatPage) Init() tea.Cmd {

@@ -106,15 +110,8 @@ func (p *chatPage) sendMessage(text string) tea.Cmd {

106 }

110 }

107 cmds = append(cmds, util.CmdHandler(chat.SessionSelectedMsg(session)))

111 cmds = append(cmds, util.CmdHandler(chat.SessionSelectedMsg(session)))

108 }

112 }

109− // TODO: move this to a service

110− a, err := agent.NewCoderAgent(p.app)

111− if err != nil {

112− return util.ReportError(err)

113− }

114− go func() {

115− a.Generate(context.Background(), p.session.ID, text)

116− }()

117

113

114+ p.app.CoderAgent.Generate(context.Background(), p.session.ID, text)

118 return tea.Batch(cmds...)

115 return tea.Batch(cmds...)

119}

116}

120

117