feat: complete v0.3.1 agent workflow

internal/agent/agent.go

@@ -25,6 +25,7 @@ const (
 	toolLogOutputPreviewBytes     = 4096
 	toolLogOutputMarker           = "[output]"
 	toolStreamLogLimit            = tools.MaxToolOutputBytes
+	continuationPrompt            = "Continue with the current task."
 
 	summarizerSystemPrompt = "You are a conversation summarizer. Summarize the following conversation history into a concise but detailed summary. Preserve:\n" +
 		"- Key facts, decisions, and conclusions\n" +
@@ -41,6 +42,7 @@ type toolExecutionResult struct {
 }
 
 type Agent struct {
+	mu               sync.RWMutex
 	provider         llm.Provider
 	providerName     string
 	model            string
@@ -49,6 +51,8 @@ type Agent struct {
 	thinkingEnabled  bool
 	systemPrompt     string
 	toolRegistry     *tools.Registry
+	authorizer       ToolAuthorizer
+	rollbackStore    *tools.RollbackStore
 	toolTimeout      time.Duration
 	maxToolRounds    int
 	maxContextTokens int
@@ -65,6 +69,8 @@ type Options struct {
 	ThinkingEnabled   bool
 	SystemPrompt      string
 	ToolRegistry      *tools.Registry
+	Authorizer        ToolAuthorizer
+	RollbackStore     *tools.RollbackStore
 	ToolTimeout       time.Duration
 	MaxToolIterations int
 	MaxContextTokens  int
@@ -88,6 +94,8 @@ func New(opts Options) *Agent {
 		thinkingEnabled:  opts.ThinkingEnabled,
 		systemPrompt:     opts.SystemPrompt,
 		toolRegistry:     opts.ToolRegistry,
+		authorizer:       opts.Authorizer,
+		rollbackStore:    opts.RollbackStore,
 		toolTimeout:      toolTimeout,
 		maxToolRounds:    maxToolRounds,
 		maxContextTokens: opts.MaxContextTokens,
@@ -116,6 +124,8 @@ type RuntimeInfo struct {
 }
 
 func (a *Agent) SetRuntime(opts RuntimeOptions) {
+	a.mu.Lock()
+	defer a.mu.Unlock()
 	if opts.Provider != nil {
 		a.provider = opts.Provider
 	}
@@ -132,6 +142,8 @@ func (a *Agent) SetRuntime(opts RuntimeOptions) {
 }
 
 func (a *Agent) RuntimeInfo() RuntimeInfo {
+	a.mu.RLock()
+	defer a.mu.RUnlock()
 	return RuntimeInfo{
 		ProviderName:     a.providerName,
 		Model:            a.model,
@@ -143,6 +155,12 @@ func (a *Agent) RuntimeInfo() RuntimeInfo {
 }
 
 func (a *Agent) Clear() {
+	a.mu.Lock()
+	defer a.mu.Unlock()
+	a.clear()
+}
+
+func (a *Agent) clear() {
 	a.lastUsage = nil
 	a.messages = nil
 	if strings.TrimSpace(a.systemPrompt) != "" {
@@ -153,8 +171,27 @@ func (a *Agent) Clear() {
 	}
 }
 
+// AddSystemInstructions makes newly installed instructions effective for the
+// current conversation and for future Clear calls.
+func (a *Agent) AddSystemInstructions(instructions string) {
+	instructions = strings.TrimSpace(instructions)
+	if instructions == "" {
+		return
+	}
+	a.mu.Lock()
+	defer a.mu.Unlock()
+	a.systemPrompt = strings.TrimSpace(a.systemPrompt + "\n\n" + instructions)
+	if len(a.messages) > 0 && a.messages[0].Role == llm.RoleSystem {
+		a.messages[0].Content = a.systemPrompt
+		return
+	}
+	a.messages = append([]llm.Message{{Role: llm.RoleSystem, Content: a.systemPrompt}}, a.messages...)
+}
+
 // Messages returns a copy of the current conversation messages.
 func (a *Agent) Messages() []llm.Message {
+	a.mu.RLock()
+	defer a.mu.RUnlock()
 	out := make([]llm.Message, len(a.messages))
 	copy(out, a.messages)
 	return out
@@ -162,11 +199,23 @@ func (a *Agent) Messages() []llm.Message {
 
 // SetMessages replaces the conversation messages (used for session restore).
 func (a *Agent) SetMessages(messages []llm.Message) {
+	a.mu.Lock()
+	defer a.mu.Unlock()
 	a.messages = append([]llm.Message(nil), messages...)
+	if strings.TrimSpace(a.systemPrompt) == "" {
+		return
+	}
+	if len(a.messages) > 0 && a.messages[0].Role == llm.RoleSystem {
+		a.messages[0].Content = a.systemPrompt
+		return
+	}
+	a.messages = append([]llm.Message{{Role: llm.RoleSystem, Content: a.systemPrompt}}, a.messages...)
 }
 
 // LastUsage returns a copy of the most recent provider token usage data.
 func (a *Agent) LastUsage() *llm.Usage {
+	a.mu.RLock()
+	defer a.mu.RUnlock()
 	if a.lastUsage == nil {
 		return nil
 	}
@@ -176,6 +225,12 @@ func (a *Agent) LastUsage() *llm.Usage {
 
 // SetLastUsage restores the most recent provider token usage data.
 func (a *Agent) SetLastUsage(usage *llm.Usage) {
+	a.mu.Lock()
+	defer a.mu.Unlock()
+	a.setLastUsage(usage)
+}
+
+func (a *Agent) setLastUsage(usage *llm.Usage) {
 	if usage == nil {
 		a.lastUsage = nil
 		return
@@ -242,27 +297,32 @@ func compactRecentStart(messages []llm.Message, prefixEnd int) int {
 // Compact summarizes older conversation messages while preserving the system
 // prompt and recent messages intact.
 func (a *Agent) Compact(ctx context.Context) error {
-	if len(a.messages) < 4 {
+	a.mu.RLock()
+	messages := append([]llm.Message(nil), a.messages...)
+	provider := a.provider
+	model := a.model
+	a.mu.RUnlock()
+	if len(messages) < 4 {
 		return nil
 	}
-	if a.provider == nil {
+	if provider == nil {
 		return fmt.Errorf("provider is not configured")
 	}
 
 	prefixEnd := 0
-	if a.messages[0].Role == llm.RoleSystem {
+	if messages[0].Role == llm.RoleSystem {
 		prefixEnd = 1
 	}
-	recentStart := compactRecentStart(a.messages, prefixEnd)
-	if recentStart <= prefixEnd || recentStart >= len(a.messages) {
+	recentStart := compactRecentStart(messages, prefixEnd)
+	if recentStart <= prefixEnd || recentStart >= len(messages) {
 		return nil
 	}
-	middle := append([]llm.Message(nil), a.messages[prefixEnd:recentStart]...)
+	middle := append([]llm.Message(nil), messages[prefixEnd:recentStart]...)
 	if len(middle) == 0 {
 		return nil
 	}
 
-	summary, usage, err := a.summarizeMessages(ctx, middle)
+	summary, usage, err := summarizeMessages(ctx, provider, model, middle)
 	if err != nil {
 		return fmt.Errorf("compact context: %w", err)
 	}
@@ -271,13 +331,15 @@ func (a *Agent) Compact(ctx context.Context) error {
 		return fmt.Errorf("compact context: summary was empty")
 	}
 
-	compressed := make([]llm.Message, 0, prefixEnd+1+len(a.messages[recentStart:]))
-	compressed = append(compressed, a.messages[:prefixEnd]...)
+	compressed := make([]llm.Message, 0, prefixEnd+1+len(messages[recentStart:]))
+	compressed = append(compressed, messages[:prefixEnd]...)
 	compressed = append(compressed, llm.Message{
 		Role:    llm.RoleAssistant,
 		Content: contextSummaryPrefix + "\n" + summary,
 	})
-	compressed = append(compressed, a.messages[recentStart:]...)
+	compressed = append(compressed, messages[recentStart:]...)
+	a.mu.Lock()
+	defer a.mu.Unlock()
 	a.messages = compressed
 	if usage != nil {
 		a.lastUsage = usage
@@ -286,6 +348,8 @@ func (a *Agent) Compact(ctx context.Context) error {
 }
 
 func (a *Agent) shouldCompact() bool {
+	a.mu.RLock()
+	defer a.mu.RUnlock()
 	if a.maxContextTokens <= 0 {
 		return false
 	}
@@ -295,9 +359,9 @@ func (a *Agent) shouldCompact() bool {
 	return llm.EstimateTokens(a.messages, a.toolDefinitions()) >= a.maxContextTokens
 }
 
-func (a *Agent) summarizeMessages(ctx context.Context, messages []llm.Message) (string, *llm.Usage, error) {
+func summarizeMessages(ctx context.Context, provider llm.Provider, model string, messages []llm.Message) (string, *llm.Usage, error) {
 	req := llm.ChatRequest{
-		Model: a.model,
+		Model: model,
 		Messages: []llm.Message{
 			{Role: llm.RoleSystem, Content: summarizerSystemPrompt},
 			{Role: llm.RoleUser, Content: formatMessagesForSummary(messages)},
@@ -306,7 +370,7 @@ func (a *Agent) summarizeMessages(ctx context.Context, messages []llm.Message) (
 
 	var summary strings.Builder
 	var usage *llm.Usage
-	if err := a.provider.ChatStream(ctx, req, func(event llm.StreamEvent) error {
+	if err := provider.ChatStream(ctx, req, func(event llm.StreamEvent) error {
 		if event.Content != "" {
 			summary.WriteString(event.Content)
 		}
@@ -354,12 +418,93 @@ func roundSignature(calls []llm.ToolCall) string {
 	return b.String()
 }
 
+// BTW answers a temporary side question using a snapshot of the current
+// conversation. It intentionally exposes no tools and does not mutate the
+// conversation, usage, or compaction state.
+func (a *Agent) BTW(ctx context.Context, input string, out io.Writer) error {
+	input = strings.TrimSpace(input)
+	if input == "" {
+		return fmt.Errorf("btw question is required")
+	}
+
+	a.mu.RLock()
+	provider := a.provider
+	model := a.model
+	messages := append([]llm.Message(nil), a.messages...)
+	thinking := a.thinking
+	thinkingParam := a.thinkingParam
+	thinkingEnabled := a.thinkingEnabled
+	a.mu.RUnlock()
+	if provider == nil {
+		return fmt.Errorf("provider is not configured")
+	}
+
+	// A snapshot taken while tools are running may end with unanswered tool
+	// calls, which providers reject. Keep only the latest complete boundary.
+	messages = completeConversationPrefix(messages)
+	messages = append(messages, llm.Message{Role: llm.RoleUser, Content: input})
+	messages = expandContinuationMessages(messages)
+	req := llm.ChatRequest{Model: model, Messages: messages}
+	if thinkingEnabled {
+		if thinkingParam == "" {
+			thinkingParam = "thinking"
+		}
+		req.Extra = map[string]any{thinkingParam: thinking}
+	}
+
+	return provider.ChatStream(ctx, req, func(event llm.StreamEvent) error {
+		if event.Content == "" || out == nil {
+			return nil
+		}
+		_, err := io.WriteString(out, event.Content)
+		return err
+	})
+}
+
+func completeConversationPrefix(messages []llm.Message) []llm.Message {
+	safeEnd := 0
+	pending := make(map[string]struct{})
+	for i, msg := range messages {
+		if len(pending) > 0 && msg.Role != llm.RoleTool {
+			break
+		}
+		switch {
+		case msg.Role == llm.RoleAssistant && len(msg.ToolCalls) > 0:
+			for _, call := range msg.ToolCalls {
+				pending[call.ID] = struct{}{}
+			}
+		case msg.Role == llm.RoleTool:
+			delete(pending, msg.ToolCallID)
+		}
+		if len(pending) == 0 {
+			safeEnd = i + 1
+		}
+	}
+	return append([]llm.Message(nil), messages[:safeEnd]...)
+}
+
+func expandContinuationMessages(messages []llm.Message) []llm.Message {
+	expanded := append([]llm.Message(nil), messages...)
+	for i := range expanded {
+		if expanded[i].Role == llm.RoleUser && strings.TrimSpace(expanded[i].Content) == "." {
+			expanded[i].Content = continuationPrompt
+		}
+	}
+	return expanded
+}
+
+func (a *Agent) appendMessage(message llm.Message) {
+	a.mu.Lock()
+	a.messages = append(a.messages, message)
+	a.mu.Unlock()
+}
+
 func (a *Agent) RunTurn(ctx context.Context, input string, out io.Writer, logs io.Writer) error {
 	input = strings.TrimSpace(input)
 	if input == "" {
 		return nil
 	}
-	a.messages = append(a.messages, llm.Message{Role: llm.RoleUser, Content: input})
+	a.appendMessage(llm.Message{Role: llm.RoleUser, Content: input})
 	if a.shouldCompact() {
 		if err := a.Compact(ctx); err != nil {
 			if logs != nil {
@@ -376,14 +521,14 @@ func (a *Agent) RunTurn(ctx context.Context, input string, out io.Writer, logs i
 			return err
 		}
 		if len(calls) == 0 {
-			a.messages = append(a.messages, llm.Message{
+			a.appendMessage(llm.Message{
 				Role:    llm.RoleAssistant,
 				Content: content,
 			})
 			return nil
 		}
 
-		a.messages = append(a.messages, llm.Message{
+		a.appendMessage(llm.Message{
 			Role:      llm.RoleAssistant,
 			Content:   content,
 			ToolCalls: calls,
@@ -418,7 +563,7 @@ func (a *Agent) RunTurn(ctx context.Context, input string, out io.Writer, logs i
 					strings.Join(toolNames, ", "),
 					doomLoopThreshold,
 				)
-				a.messages = append(a.messages, llm.Message{
+				a.appendMessage(llm.Message{
 					Role:    llm.RoleAssistant,
 					Content: msg,
 				})
@@ -428,7 +573,7 @@ func (a *Agent) RunTurn(ctx context.Context, input string, out io.Writer, logs i
 
 		results := a.executeTools(ctx, calls, logs)
 		for _, result := range results {
-			a.messages = append(a.messages, llm.Message{
+			a.appendMessage(llm.Message{
 				Role:       llm.RoleTool,
 				ToolCallID: result.call.ID,
 				Content:    result.output,
@@ -444,23 +589,33 @@ func (a *Agent) streamAssistant(ctx context.Context, out io.Writer) (string, []l
 	var content strings.Builder
 	calls := make(map[int]*toolCallBuilder)
 
+	a.mu.RLock()
+	provider := a.provider
+	model := a.model
+	messages := append([]llm.Message(nil), a.messages...)
+	thinking := a.thinking
+	thinkingParam := a.thinkingParam
+	thinkingEnabled := a.thinkingEnabled
+	a.mu.RUnlock()
+	if provider == nil {
+		return "", nil, fmt.Errorf("provider is not configured")
+	}
 	req := llm.ChatRequest{
-		Model:    a.model,
-		Messages: append([]llm.Message(nil), a.messages...),
+		Model:    model,
+		Messages: expandContinuationMessages(messages),
 		Tools:    a.toolDefinitions(),
 	}
-	if a.thinkingEnabled {
-		thinkingParam := a.thinkingParam
+	if thinkingEnabled {
 		if thinkingParam == "" {
 			thinkingParam = "thinking"
 		}
-		req.Extra = map[string]any{thinkingParam: a.thinking}
+		req.Extra = map[string]any{thinkingParam: thinking}
 	}
 	if len(req.Tools) > 0 {
 		req.ToolChoice = "auto"
 	}
 
-	err := a.provider.ChatStream(ctx, req, func(event llm.StreamEvent) error {
+	err := provider.ChatStream(ctx, req, func(event llm.StreamEvent) error {
 		if event.Usage != nil {
 			a.SetLastUsage(event.Usage)
 		}
@@ -513,21 +668,90 @@ func (a *Agent) executeTool(ctx context.Context, call llm.ToolCall, logs io.Writ
 	}
 	raw := json.RawMessage(args)
 
+	rollbackNotes, err := a.prepareMutation(ctx, tool, call, raw, args)
+	if err != nil {
+		return tools.FormatError(err), false
+	}
+
 	toolCtx, cancel := context.WithTimeout(ctx, a.toolTimeout)
 	defer cancel()
 
 	if streaming, ok := tool.(tools.StreamingTool); ok {
-		return a.executeStreamingTool(streaming, toolCtx, raw, name, args, logs)
+		output, streamed := a.executeStreamingTool(streaming, toolCtx, raw, name, args, logs)
+		return appendRollbackNotes(output, rollbackNotes), streamed
 	}
 
 	output, err := tool.Execute(toolCtx, raw)
 	if err != nil {
 		if output != "" {
-			return output + "\n" + tools.FormatError(err), false
+			return appendRollbackNotes(output+"\n"+tools.FormatError(err), rollbackNotes), false
 		}
-		return tools.FormatError(err), false
+		return appendRollbackNotes(tools.FormatError(err), rollbackNotes), false
 	}
-	return output, false
+	return appendRollbackNotes(output, rollbackNotes), false
+}
+
+func (a *Agent) prepareMutation(ctx context.Context, tool tools.Tool, call llm.ToolCall, raw json.RawMessage, args string) ([]string, error) {
+	if !tools.IsMutatingCall(tool, raw) {
+		return nil, nil
+	}
+
+	var targets []string
+	if targetTool, ok := tool.(tools.TargetPathTool); ok {
+		var err error
+		targets, err = targetTool.TargetPaths(raw)
+		if err != nil {
+			return nil, fmt.Errorf("resolve mutation targets: %w", err)
+		}
+	}
+
+	preview := compact(args, toolLogArgumentLimit)
+	if previewTool, ok := tool.(tools.PreviewableTool); ok {
+		value, err := previewTool.Preview(ctx, raw)
+		if err != nil {
+			return nil, fmt.Errorf("preview mutation: %w", err)
+		}
+		preview = value
+	}
+
+	if a.authorizer != nil {
+		decision, err := a.authorizer.AuthorizeToolCall(ctx, ToolAuthorizationRequest{
+			ID:                call.ID,
+			ToolName:          call.Function.Name,
+			Arguments:         args,
+			Preview:           preview,
+			RollbackAvailable: a.rollbackStore != nil && len(targets) > 0,
+		})
+		if err != nil {
+			return nil, fmt.Errorf("authorize %s: %w", call.Function.Name, err)
+		}
+		if decision != ToolAuthorizationAllow {
+			return nil, fmt.Errorf("tool call denied by user: %s", call.Function.Name)
+		}
+	}
+
+	if a.rollbackStore == nil {
+		return nil, nil
+	}
+	notes := make([]string, 0, len(targets))
+	for _, path := range targets {
+		id, err := a.rollbackStore.Capture(call.Function.Name, path, path)
+		if err != nil {
+			return nil, fmt.Errorf("capture rollback snapshot for %s: %w", path, err)
+		}
+		notes = append(notes, fmt.Sprintf("[rollback] %s %s", id, path))
+	}
+	return notes, nil
+}
+
+func appendRollbackNotes(output string, notes []string) string {
+	if len(notes) == 0 {
+		return output
+	}
+	if strings.TrimSpace(output) == "" {
+		return strings.Join(notes, "\n")
+	}
+	return output + "\n" + strings.Join(notes, "\n")
 }
 
 func (a *Agent) executeStreamingTool(streaming tools.StreamingTool, ctx context.Context, raw json.RawMessage, name, args string, logs io.Writer) (string, bool) {
@@ -573,7 +797,11 @@ func (a *Agent) hasMutatingCall(calls []llm.ToolCall) bool {
 	}
 	for _, call := range calls {
 		if tool, ok := a.toolRegistry.Get(call.Function.Name); ok {
-			if tools.IsMutating(tool) {
+			args := strings.TrimSpace(call.Function.Arguments)
+			if args == "" {
+				args = "{}"
+			}
+			if tools.IsMutatingCall(tool, json.RawMessage(args)) {
 				return true
 			}
 		}