chore: cleanup

4 files changed, 36 insertions, 77 deletions

diff --git a/internal/llm/agent/agent-tool.go b/internal/llm/agent/agent-tool.go
index d451eafed..85ec4b2a4 100644
--- a/internal/llm/agent/agent-tool.go
+++ b/internal/llm/agent/agent-tool.go
@@ -88,8 +88,6 @@ func (b *agentTool) Run(ctx context.Context, call tools.ToolCall) (tools.ToolRes
 	}
 
 	parentSession.Cost += updatedSession.Cost
-	parentSession.PromptTokens += updatedSession.PromptTokens
-	parentSession.CompletionTokens += updatedSession.CompletionTokens
 
 	_, err = b.sessions.Update(ctx, parentSession)
 	if err != nil {
diff --git a/internal/llm/agent/agent.go b/internal/llm/agent/agent.go
index f6f816256..778c0bcf4 100644
--- a/internal/llm/agent/agent.go
+++ b/internal/llm/agent/agent.go
@@ -300,53 +300,6 @@ func (a *agent) createUserMessage(ctx context.Context, sessionID, content string
 }
 
 func (a *agent) streamAndHandleEvents(ctx context.Context, sessionID string, msgHistory []message.Message) (message.Message, *message.Message, error) {
-	// Check if we need to auto-compact based on token count
-	contextWindow := a.provider.Model().ContextWindow
-	maxTokens := a.provider.MaxTokens()
-	threshold := int64(float64(contextWindow) * 0.80)
-	usage, err := a.GetUsage(ctx, sessionID)
-	if err != nil || usage == nil {
-		return message.Message{}, nil, fmt.Errorf("failed to get usage: %w", err)
-	}
-
-	// If we're approaching the context window limit, trigger auto-compaction
-	if false && (*usage+maxTokens) >= threshold {
-		status.Info(fmt.Sprintf("Auto-compaction triggered for session %s. Estimated tokens: %d, Threshold: %d", sessionID, usage, threshold))
-
-		// Perform compaction with pause/resume to ensure safety
-		if err := a.CompactSession(ctx, sessionID); err != nil {
-			status.Error(fmt.Sprintf("Auto-compaction failed: %v", err))
-			// Continue with the request even if compaction fails
-		} else {
-			// Re-fetch session details after compaction
-			currentSession, err := a.sessions.Get(ctx, sessionID)
-			if err != nil {
-				return message.Message{}, nil, fmt.Errorf("failed to get session after compaction: %w", err)
-			}
-
-			// Re-prepare messages using the new summary
-			var sessionMessages []message.Message
-			if currentSession.Summary != "" && currentSession.SummarizedAt > 0 {
-				// If summary exists, only fetch messages after the summarization timestamp
-				sessionMessages, err = a.messages.ListAfter(ctx, sessionID, currentSession.SummarizedAt)
-				if err != nil {
-					return message.Message{}, nil, fmt.Errorf("failed to list messages after compaction: %w", err)
-				}
-
-				// Create a new message history with the summary and messages after summarization
-				summaryMessage := message.Message{
-					Role: message.Assistant,
-					Parts: []message.ContentPart{
-						message.TextContent{Text: currentSession.Summary},
-					},
-				}
-
-				// Replace msgHistory with the new compacted version
-				msgHistory = append([]message.Message{summaryMessage}, sessionMessages...)
-			}
-		}
-	}
-
 	eventChan := a.provider.StreamResponse(ctx, msgHistory, a.tools)
 
 	assistantMsg, err := a.messages.Create(ctx, sessionID, message.CreateMessageParams{
@@ -541,8 +494,8 @@ func (a *agent) TrackUsage(ctx context.Context, sessionID string, model models.M
 		model.CostPer1MOut/1e6*float64(usage.OutputTokens)
 
 	sess.Cost += cost
-	sess.CompletionTokens += usage.OutputTokens
-	sess.PromptTokens += usage.InputTokens
+	sess.CompletionTokens = usage.OutputTokens + usage.CacheReadTokens
+	sess.PromptTokens = usage.InputTokens + usage.CacheCreationTokens
 
 	_, err = a.sessions.Update(ctx, sess)
 	if err != nil {
@@ -646,7 +599,16 @@ func (a *agent) CompactSession(ctx context.Context, sessionID string) error {
 			Role: message.System,
 			Parts: []message.ContentPart{
 				message.TextContent{
-					Text: "You are a helpful AI assistant tasked with summarizing conversations.",
+					Text: `You are a helpful AI assistant tasked with summarizing conversations.
+
+When asked to summarize, provide a detailed but concise summary of the conversation. 
+Focus on information that would be helpful for continuing the conversation, including:
+- What was done
+- What is currently being worked on
+- Which files are being modified
+- What needs to be done next
+
+Your summary should be comprehensive enough to provide context but concise enough to be quickly understood.`,
 				},
 			},
 		},
@@ -655,7 +617,7 @@ func (a *agent) CompactSession(ctx context.Context, sessionID string) error {
 	// If there's an existing summary, include it
 	if existingSummary != "" {
 		messages = append(messages, message.Message{
-			Role: message.Assistant, // TODO: should this be system or user instead?
+			Role: message.Assistant,
 			Parts: []message.ContentPart{
 				message.TextContent{
 					Text: existingSummary,
diff --git a/internal/llm/provider/bedrock.go b/internal/llm/provider/bedrock.go
index 9f42e5b18..9fa3ca87f 100644
--- a/internal/llm/provider/bedrock.go
+++ b/internal/llm/provider/bedrock.go
@@ -98,4 +98,3 @@ func (b *bedrockClient) stream(ctx context.Context, messages []message.Message,
 
 	return b.childProvider.stream(ctx, messages, tools)
 }
-
diff --git a/internal/llm/tools/ls_test.go b/internal/llm/tools/ls_test.go
index 508cb98d3..07003c44a 100644
--- a/internal/llm/tools/ls_test.go
+++ b/internal/llm/tools/ls_test.go
@@ -83,19 +83,19 @@ func TestLsTool_Run(t *testing.T) {
 
 		response, err := tool.Run(context.Background(), call)
 		require.NoError(t, err)
-		
+
 		// Check that visible directories and files are included
 		assert.Contains(t, response.Content, "dir1")
 		assert.Contains(t, response.Content, "dir2")
 		assert.Contains(t, response.Content, "dir3")
 		assert.Contains(t, response.Content, "file1.txt")
 		assert.Contains(t, response.Content, "file2.txt")
-		
+
 		// Check that hidden files and directories are not included
 		assert.NotContains(t, response.Content, ".hidden_dir")
 		assert.NotContains(t, response.Content, ".hidden_file.txt")
 		assert.NotContains(t, response.Content, ".hidden_root_file.txt")
-		
+
 		// Check that __pycache__ is not included
 		assert.NotContains(t, response.Content, "__pycache__")
 	})
@@ -122,7 +122,7 @@ func TestLsTool_Run(t *testing.T) {
 	t.Run("handles empty path parameter", func(t *testing.T) {
 		// For this test, we need to mock the config.WorkingDirectory function
 		// Since we can't easily do that, we'll just check that the response doesn't contain an error message
-		
+
 		tool := NewLsTool()
 		params := LSParams{
 			Path: "",
@@ -138,7 +138,7 @@ func TestLsTool_Run(t *testing.T) {
 
 		response, err := tool.Run(context.Background(), call)
 		require.NoError(t, err)
-		
+
 		// The response should either contain a valid directory listing or an error
 		// We'll just check that it's not empty
 		assert.NotEmpty(t, response.Content)
@@ -173,11 +173,11 @@ func TestLsTool_Run(t *testing.T) {
 
 		response, err := tool.Run(context.Background(), call)
 		require.NoError(t, err)
-		
+
 		// The output format is a tree, so we need to check for specific patterns
 		// Check that file1.txt is not directly mentioned
 		assert.NotContains(t, response.Content, "- file1.txt")
-		
+
 		// Check that dir1/ is not directly mentioned
 		assert.NotContains(t, response.Content, "- dir1/")
 	})
@@ -189,12 +189,12 @@ func TestLsTool_Run(t *testing.T) {
 		defer func() {
 			os.Chdir(origWd)
 		}()
-		
+
 		// Change to a directory above the temp directory
 		parentDir := filepath.Dir(tempDir)
 		err = os.Chdir(parentDir)
 		require.NoError(t, err)
-		
+
 		tool := NewLsTool()
 		params := LSParams{
 			Path: filepath.Base(tempDir),
@@ -210,7 +210,7 @@ func TestLsTool_Run(t *testing.T) {
 
 		response, err := tool.Run(context.Background(), call)
 		require.NoError(t, err)
-		
+
 		// Should list the temp directory contents
 		assert.Contains(t, response.Content, "dir1")
 		assert.Contains(t, response.Content, "file1.txt")
@@ -291,22 +291,22 @@ func TestCreateFileTree(t *testing.T) {
 	}
 
 	tree := createFileTree(paths)
-	
+
 	// Check the structure of the tree
 	assert.Len(t, tree, 1) // Should have one root node
-	
+
 	// Check the root node
 	rootNode := tree[0]
 	assert.Equal(t, "path", rootNode.Name)
 	assert.Equal(t, "directory", rootNode.Type)
 	assert.Len(t, rootNode.Children, 1)
-	
+
 	// Check the "to" node
 	toNode := rootNode.Children[0]
 	assert.Equal(t, "to", toNode.Name)
 	assert.Equal(t, "directory", toNode.Type)
 	assert.Len(t, toNode.Children, 3) // file1.txt, dir1, dir2
-	
+
 	// Find the dir1 node
 	var dir1Node *TreeNode
 	for _, child := range toNode.Children {
@@ -315,7 +315,7 @@ func TestCreateFileTree(t *testing.T) {
 			break
 		}
 	}
-	
+
 	require.NotNil(t, dir1Node)
 	assert.Equal(t, "directory", dir1Node.Type)
 	assert.Len(t, dir1Node.Children, 2) // file2.txt and subdir
@@ -354,9 +354,9 @@ func TestPrintTree(t *testing.T) {
 			Type: "file",
 		},
 	}
-	
+
 	result := printTree(tree, "/root")
-	
+
 	// Check the output format
 	assert.Contains(t, result, "- /root/")
 	assert.Contains(t, result, "  - dir1/")
@@ -405,7 +405,7 @@ func TestListDirectory(t *testing.T) {
 		files, truncated, err := listDirectory(tempDir, []string{}, 1000)
 		require.NoError(t, err)
 		assert.False(t, truncated)
-		
+
 		// Check that visible files and directories are included
 		containsPath := func(paths []string, target string) bool {
 			targetPath := filepath.Join(tempDir, target)
@@ -416,12 +416,12 @@ func TestListDirectory(t *testing.T) {
 			}
 			return false
 		}
-		
+
 		assert.True(t, containsPath(files, "dir1"))
 		assert.True(t, containsPath(files, "file1.txt"))
 		assert.True(t, containsPath(files, "file2.txt"))
 		assert.True(t, containsPath(files, "dir1/file3.txt"))
-		
+
 		// Check that hidden files and directories are not included
 		assert.False(t, containsPath(files, ".hidden_dir"))
 		assert.False(t, containsPath(files, ".hidden_file.txt"))
@@ -438,12 +438,12 @@ func TestListDirectory(t *testing.T) {
 		files, truncated, err := listDirectory(tempDir, []string{"*.txt"}, 1000)
 		require.NoError(t, err)
 		assert.False(t, truncated)
-		
+
 		// Check that no .txt files are included
 		for _, file := range files {
 			assert.False(t, strings.HasSuffix(file, ".txt"), "Found .txt file: %s", file)
 		}
-		
+
 		// But directories should still be included
 		containsDir := false
 		for _, file := range files {
@@ -454,4 +454,4 @@ func TestListDirectory(t *testing.T) {
 		}
 		assert.True(t, containsDir)
 	})
-}
\ No newline at end of file
+}