Plexum-gemini-provider/internal/runner/session_mutate_apply.go

// session_mutate_apply — gemini consume mirror. Rewrites the chat
// JSONL so the consumed view propagates to next-resume gemini.
//
// Per-line mutation:
//
//	type="gemini" line: every toolCalls[i].id matching → if heavy,
//	  reset args to {}
//	type="user" line:  every content[j].functionResponse.id matching
//	  → set response.output to the marker

package runner

import (
	"bufio"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"strings"
)

const (
	consumedMarker   = "...(consumed)"
	toolCalledMarker = "...(tool called)"
)

// GeminiMutation mirrors plugin.BlockMutation.
type GeminiMutation struct {
	BlockID string
	Op      string
}

// MutateGeminiSession edits gemini's chat JSONL in lockstep with
// Plexum's dynamic.jsonl. Returns the count of distinct ids touched.
func MutateGeminiSession(workspace string, mutations []GeminiMutation) (int, error) {
	if len(mutations) == 0 {
		return 0, nil
	}
	path, err := FindGeminiSessionFile(workspace)
	if err != nil {
		return 0, nil // no session yet
	}
	byID := make(map[string]string, len(mutations))
	for _, m := range mutations {
		byID[m.BlockID] = m.Op
	}
	f, err := os.Open(path)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			return 0, nil
		}
		return 0, fmt.Errorf("open gemini session: %w", err)
	}
	defer f.Close()

	var out strings.Builder
	touched := map[string]struct{}{}
	mutated := false
	sc := bufio.NewScanner(f)
	sc.Buffer(make([]byte, 64*1024), 16*1024*1024)
	for sc.Scan() {
		line := sc.Bytes()
		if len(line) == 0 {
			out.WriteByte('\n')
			continue
		}
		rewritten, hits := rewriteGeminiLine(line, byID)
		out.Write(rewritten)
		out.WriteByte('\n')
		for _, h := range hits {
			touched[h] = struct{}{}
			mutated = true
		}
	}
	if err := sc.Err(); err != nil {
		return 0, fmt.Errorf("scan gemini session: %w", err)
	}
	if !mutated {
		return 0, nil
	}
	if err := writeAtomic(path, []byte(out.String())); err != nil {
		return len(touched), fmt.Errorf("rewrite gemini session: %w", err)
	}
	return len(touched), nil
}

// rewriteGeminiLine returns rewritten bytes + list of touched ids.
func rewriteGeminiLine(line []byte, byID map[string]string) ([]byte, []string) {
	var obj map[string]any
	if err := json.Unmarshal(line, &obj); err != nil {
		return line, nil
	}
	var touched []string
	mutated := false

	switch obj["type"] {
	case "gemini":
		tcs, ok := obj["toolCalls"].([]any)
		if !ok {
			return line, nil
		}
		for i, tcAny := range tcs {
			tc, ok := tcAny.(map[string]any)
			if !ok {
				continue
			}
			id, _ := tc["id"].(string)
			op, ok := byID[id]
			if !ok || id == "" {
				continue
			}
			if op == "consume-heavy" {
				tc["args"] = map[string]any{}
				tcs[i] = tc
				touched = append(touched, id)
				mutated = true
			}
		}
		obj["toolCalls"] = tcs
	case "user":
		contents, ok := obj["content"].([]any)
		if !ok {
			return line, nil
		}
		for i, ctAny := range contents {
			ct, ok := ctAny.(map[string]any)
			if !ok {
				continue
			}
			fr, ok := ct["functionResponse"].(map[string]any)
			if !ok {
				continue
			}
			id, _ := fr["id"].(string)
			op, ok := byID[id]
			if !ok || id == "" {
				continue
			}
			marker := consumedMarker
			if op == "consume-heavy" {
				marker = toolCalledMarker
			}
			resp, _ := fr["response"].(map[string]any)
			if resp == nil {
				resp = map[string]any{}
			}
			resp["output"] = marker
			fr["response"] = resp
			ct["functionResponse"] = fr
			contents[i] = ct
			touched = append(touched, id)
			mutated = true
		}
		obj["content"] = contents
	default:
		return line, nil
	}
	if !mutated {
		return line, nil
	}
	rewritten, err := json.Marshal(obj)
	if err != nil {
		return line, nil
	}
	return rewritten, touched
}

func writeAtomic(path string, data []byte) error {
	tmp := path + ".tmp"
	if err := os.WriteFile(tmp, data, 0o600); err != nil {
		return err
	}
	return os.Rename(tmp, path)
}

// FindGeminiSessionFile resolves ~/.gemini/tmp/<ws>/chats/session-*-<sid_prefix>.jsonl
// for the captured session id. Picks the most-recent match by mtime
// since gemini may roll the chat across multiple files within the
// same logical session id.
func FindGeminiSessionFile(workspace string) (string, error) {
	sid := loadSessionID(workspace)
	if sid == "" {
		return "", errors.New("no gemini session id captured yet")
	}
	prefix := sid
	if len(prefix) > 8 {
		prefix = prefix[:8]
	}
	home, err := os.UserHomeDir()
	if err != nil {
		return "", err
	}
	wsName := filepathBase(workspace)
	dir := home + "/.gemini/tmp/" + wsName + "/chats"
	entries, err := os.ReadDir(dir)
	if err != nil {
		return "", fmt.Errorf("read gemini chats dir %s: %w", dir, err)
	}
	type cand struct {
		path  string
		mtime int64
	}
	var matches []cand
	for _, e := range entries {
		if e.IsDir() || !strings.HasSuffix(e.Name(), ".jsonl") {
			continue
		}
		if !strings.Contains(e.Name(), prefix) {
			continue
		}
		info, err := e.Info()
		if err != nil {
			continue
		}
		matches = append(matches, cand{path: dir + "/" + e.Name(), mtime: info.ModTime().UnixNano()})
	}
	if len(matches) == 0 {
		return "", fmt.Errorf("no chat file matching session_id prefix %q in %s", prefix, dir)
	}
	// pick newest
	best := matches[0]
	for _, m := range matches[1:] {
		if m.mtime > best.mtime {
			best = m
		}
	}
	return best.path, nil
}

func filepathBase(p string) string {
	// minimal local impl to keep imports tight
	i := strings.LastIndex(p, "/")
	if i < 0 {
		return p
	}
	return p[i+1:]
}