wip: refactoring

This commit is contained in:
adamdottv
2025-05-12 09:44:56 -05:00
parent ed9fba99c9
commit dfe5fd8d97
19 changed files with 97 additions and 75 deletions

View File

@@ -91,7 +91,7 @@ func (b *agentTool) Run(ctx context.Context, call tools.ToolCall) (tools.ToolRes
parentSession.PromptTokens += updatedSession.PromptTokens
parentSession.CompletionTokens += updatedSession.CompletionTokens
_, err = b.sessions.Save(ctx, parentSession)
_, err = b.sessions.Update(ctx, parentSession)
if err != nil {
return tools.ToolResponse{}, fmt.Errorf("error saving parent session: %s", err)
}

View File

@@ -156,7 +156,7 @@ func (a *agent) generateTitle(ctx context.Context, sessionID string, content str
}
session.Title = title
_, err = a.sessions.Save(ctx, session)
_, err = a.sessions.Update(ctx, session)
return err
}
@@ -459,7 +459,7 @@ out:
func (a *agent) finishMessage(ctx context.Context, msg *message.Message, finishReson message.FinishReason) {
msg.AddFinish(finishReson)
_ = a.messages.Update(ctx, *msg)
_, _ = a.messages.Update(ctx, *msg)
}
func (a *agent) processEvent(ctx context.Context, sessionID string, assistantMsg *message.Message, event provider.ProviderEvent) error {
@@ -477,13 +477,16 @@ func (a *agent) processEvent(ctx context.Context, sessionID string, assistantMsg
switch event.Type {
case provider.EventThinkingDelta:
assistantMsg.AppendReasoningContent(event.Content)
return a.messages.Update(ctx, *assistantMsg)
_, err := a.messages.Update(ctx, *assistantMsg)
return err
case provider.EventContentDelta:
assistantMsg.AppendContent(event.Content)
return a.messages.Update(ctx, *assistantMsg)
_, err := a.messages.Update(ctx, *assistantMsg)
return err
case provider.EventToolUseStart:
assistantMsg.AddToolCall(*event.ToolCall)
return a.messages.Update(ctx, *assistantMsg)
_, err := a.messages.Update(ctx, *assistantMsg)
return err
// TODO: see how to handle this
// case provider.EventToolUseDelta:
// tm := time.Unix(assistantMsg.UpdatedAt, 0)
@@ -495,7 +498,8 @@ func (a *agent) processEvent(ctx context.Context, sessionID string, assistantMsg
// }
case provider.EventToolUseStop:
assistantMsg.FinishToolCall(event.ToolCall.ID)
return a.messages.Update(ctx, *assistantMsg)
_, err := a.messages.Update(ctx, *assistantMsg)
return err
case provider.EventError:
if errors.Is(event.Error, context.Canceled) {
status.Info(fmt.Sprintf("Event processing canceled for session: %s", sessionID))
@@ -506,7 +510,7 @@ func (a *agent) processEvent(ctx context.Context, sessionID string, assistantMsg
case provider.EventComplete:
assistantMsg.SetToolCalls(event.Response.ToolCalls)
assistantMsg.AddFinish(event.Response.FinishReason)
if err := a.messages.Update(ctx, *assistantMsg); err != nil {
if _, err := a.messages.Update(ctx, *assistantMsg); err != nil {
return fmt.Errorf("failed to update message: %w", err)
}
return a.TrackUsage(ctx, sessionID, a.provider.Model(), event.Response.Usage)
@@ -540,7 +544,7 @@ func (a *agent) TrackUsage(ctx context.Context, sessionID string, model models.M
sess.CompletionTokens += usage.OutputTokens
sess.PromptTokens += usage.InputTokens
_, err = a.sessions.Save(ctx, sess)
_, err = a.sessions.Update(ctx, sess)
if err != nil {
return fmt.Errorf("failed to save session: %w", err)
}
@@ -691,7 +695,7 @@ func (a *agent) CompactSession(ctx context.Context, sessionID string) error {
session.SummarizedAt = currentTime
// Save the updated session
_, err = a.sessions.Save(ctx, session)
_, err = a.sessions.Update(ctx, session)
if err != nil {
return fmt.Errorf("failed to save session with summary: %w", err)
}

View File

@@ -86,6 +86,7 @@ func (b *mcpTool) Run(ctx context.Context, params tools.ToolCall) (tools.ToolRes
}
permissionDescription := fmt.Sprintf("execute %s with the following parameters: %s", b.Info().Name, params.Input)
p := b.permissions.Request(
ctx,
permission.CreatePermissionRequest{
SessionID: sessionID,
Path: config.WorkingDirectory(),

View File

@@ -268,6 +268,7 @@ func (b *bashTool) Run(ctx context.Context, call ToolCall) (ToolResponse, error)
}
if !isSafeReadOnly {
p := b.permissions.Request(
ctx,
permission.CreatePermissionRequest{
SessionID: sessionID,
Path: config.WorkingDirectory(),

View File

@@ -37,7 +37,7 @@ type EditResponseMetadata struct {
type editTool struct {
lspClients map[string]*lsp.Client
permissions permission.Service
files history.Service
history history.Service
}
const (
@@ -95,7 +95,7 @@ func NewEditTool(lspClients map[string]*lsp.Client, permissions permission.Servi
return &editTool{
lspClients: lspClients,
permissions: permissions,
files: files,
history: files,
}
}
@@ -202,6 +202,7 @@ func (e *editTool) createNewFile(ctx context.Context, filePath, content string)
permissionPath = rootDir
}
p := e.permissions.Request(
ctx,
permission.CreatePermissionRequest{
SessionID: sessionID,
Path: permissionPath,
@@ -224,14 +225,14 @@ func (e *editTool) createNewFile(ctx context.Context, filePath, content string)
}
// File can't be in the history so we create a new file history
_, err = e.files.Create(ctx, sessionID, filePath, "")
_, err = e.history.Create(ctx, sessionID, filePath, "")
if err != nil {
// Log error but don't fail the operation
return ToolResponse{}, fmt.Errorf("error creating file history: %w", err)
}
// Add the new content to the file history
_, err = e.files.CreateVersion(ctx, sessionID, filePath, content)
_, err = e.history.CreateVersion(ctx, sessionID, filePath, content)
if err != nil {
// Log error but don't fail the operation
slog.Debug("Error creating file history version", "error", err)
@@ -313,6 +314,7 @@ func (e *editTool) deleteContent(ctx context.Context, filePath, oldString string
permissionPath = rootDir
}
p := e.permissions.Request(
ctx,
permission.CreatePermissionRequest{
SessionID: sessionID,
Path: permissionPath,
@@ -335,9 +337,9 @@ func (e *editTool) deleteContent(ctx context.Context, filePath, oldString string
}
// Check if file exists in history
file, err := e.files.GetByPathAndSession(ctx, filePath, sessionID)
file, err := e.history.GetLatestByPathAndSession(ctx, filePath, sessionID)
if err != nil {
_, err = e.files.Create(ctx, sessionID, filePath, oldContent)
_, err = e.history.Create(ctx, sessionID, filePath, oldContent)
if err != nil {
// Log error but don't fail the operation
return ToolResponse{}, fmt.Errorf("error creating file history: %w", err)
@@ -345,13 +347,13 @@ func (e *editTool) deleteContent(ctx context.Context, filePath, oldString string
}
if file.Content != oldContent {
// User Manually changed the content store an intermediate version
_, err = e.files.CreateVersion(ctx, sessionID, filePath, oldContent)
_, err = e.history.CreateVersion(ctx, sessionID, filePath, oldContent)
if err != nil {
slog.Debug("Error creating file history version", "error", err)
}
}
// Store the new version
_, err = e.files.CreateVersion(ctx, sessionID, filePath, "")
_, err = e.history.CreateVersion(ctx, sessionID, filePath, "")
if err != nil {
slog.Debug("Error creating file history version", "error", err)
}
@@ -433,6 +435,7 @@ func (e *editTool) replaceContent(ctx context.Context, filePath, oldString, newS
permissionPath = rootDir
}
p := e.permissions.Request(
ctx,
permission.CreatePermissionRequest{
SessionID: sessionID,
Path: permissionPath,
@@ -455,9 +458,9 @@ func (e *editTool) replaceContent(ctx context.Context, filePath, oldString, newS
}
// Check if file exists in history
file, err := e.files.GetByPathAndSession(ctx, filePath, sessionID)
file, err := e.history.GetLatestByPathAndSession(ctx, filePath, sessionID)
if err != nil {
_, err = e.files.Create(ctx, sessionID, filePath, oldContent)
_, err = e.history.Create(ctx, sessionID, filePath, oldContent)
if err != nil {
// Log error but don't fail the operation
return ToolResponse{}, fmt.Errorf("error creating file history: %w", err)
@@ -465,13 +468,13 @@ func (e *editTool) replaceContent(ctx context.Context, filePath, oldString, newS
}
if file.Content != oldContent {
// User Manually changed the content store an intermediate version
_, err = e.files.CreateVersion(ctx, sessionID, filePath, oldContent)
_, err = e.history.CreateVersion(ctx, sessionID, filePath, oldContent)
if err != nil {
slog.Debug("Error creating file history version", "error", err)
}
}
// Store the new version
_, err = e.files.CreateVersion(ctx, sessionID, filePath, newContent)
_, err = e.history.CreateVersion(ctx, sessionID, filePath, newContent)
if err != nil {
slog.Debug("Error creating file history version", "error", err)
}

View File

@@ -122,6 +122,7 @@ func (t *fetchTool) Run(ctx context.Context, call ToolCall) (ToolResponse, error
}
p := t.permissions.Request(
ctx,
permission.CreatePermissionRequest{
SessionID: sessionID,
Path: config.WorkingDirectory(),

View File

@@ -193,6 +193,7 @@ func (p *patchTool) Run(ctx context.Context, call ToolCall) (ToolResponse, error
dir := filepath.Dir(path)
patchDiff, _, _ := diff.GenerateDiff("", *change.NewContent, path)
p := p.permissions.Request(
ctx,
permission.CreatePermissionRequest{
SessionID: sessionID,
Path: dir,
@@ -220,6 +221,7 @@ func (p *patchTool) Run(ctx context.Context, call ToolCall) (ToolResponse, error
patchDiff, _, _ := diff.GenerateDiff(currentContent, newContent, path)
dir := filepath.Dir(path)
p := p.permissions.Request(
ctx,
permission.CreatePermissionRequest{
SessionID: sessionID,
Path: dir,
@@ -239,6 +241,7 @@ func (p *patchTool) Run(ctx context.Context, call ToolCall) (ToolResponse, error
dir := filepath.Dir(path)
patchDiff, _, _ := diff.GenerateDiff(*change.OldContent, "", path)
p := p.permissions.Request(
ctx,
permission.CreatePermissionRequest{
SessionID: sessionID,
Path: dir,
@@ -313,7 +316,7 @@ func (p *patchTool) Run(ctx context.Context, call ToolCall) (ToolResponse, error
totalRemovals += removals
// Update history
file, err := p.files.GetByPathAndSession(ctx, absPath, sessionID)
file, err := p.files.GetLatestByPathAndSession(ctx, absPath, sessionID)
if err != nil && change.Type != diff.ActionAdd {
// If not adding a file, create history entry for existing file
_, err = p.files.Create(ctx, sessionID, absPath, oldContent)

View File

@@ -167,6 +167,7 @@ func (w *writeTool) Run(ctx context.Context, call ToolCall) (ToolResponse, error
permissionPath = rootDir
}
p := w.permissions.Request(
ctx,
permission.CreatePermissionRequest{
SessionID: sessionID,
Path: permissionPath,
@@ -189,7 +190,7 @@ func (w *writeTool) Run(ctx context.Context, call ToolCall) (ToolResponse, error
}
// Check if file exists in history
file, err := w.files.GetByPathAndSession(ctx, filePath, sessionID)
file, err := w.files.GetLatestByPathAndSession(ctx, filePath, sessionID)
if err != nil {
_, err = w.files.Create(ctx, sessionID, filePath, oldContent)
if err != nil {