X-Git-Url: https://git.giacomolli.uk/onechat/blobdiff_plain/522d7fd117285e4e4ce3ea5ec9b5a1722521b681..refs/heads/master:/whatsapp_import.go?ds=inline

diff --git a/whatsapp_import.go b/whatsapp_import.go
index 81eb808..1092b49 100644
--- a/whatsapp_import.go
+++ b/whatsapp_import.go
@@ -37,21 +37,12 @@ var (
 // whatsappImport runs the offline import command. It holds the same token lock
 // as the daemon and probe, even for a plan. The plan reads no credentials for a
 // model and makes no network call. Only a completed rebuild replaces the agent
-// directory, and unfinished subagents prevent either mode.
+// directory. Unfinished subagents and pending target deliveries prevent either
+// mode.
 func whatsappImport(ctx context.Context, args []string, out io.Writer) (err error) {
-	planOnly := len(args) > 0 && args[0] == "-n"
-	if planOnly {
-		args = args[1:]
-	}
-
-	if len(args) == 0 {
-		return errors.New("usage: onechat import [-n] ZIP [ZIP...]")
-	}
-
-	for _, arg := range args {
-		if strings.HasPrefix(arg, "-") {
-			return errors.New("usage: onechat import [-n] ZIP [ZIP...]")
-		}
+	options, err := whatsappImportParse(args)
+	if err != nil {
+		return fmt.Errorf("parse import arguments: %w", err)
 	}
 
 	dir, err := dataDir()
@@ -72,47 +63,57 @@ func whatsappImport(ctx context.Context, args []string, out io.Writer) (err erro
 	}()
 
 	root := filepath.Join(dir, "agent")
-	err = rebuildIdle(root)
+	agent, err := whatsappImportAgent(root, options.agent)
+	if err != nil {
+		return fmt.Errorf("select import agent: %w", err)
+	}
+
+	err = whatsappImportIdle(root, agent)
 	if err != nil {
 		return fmt.Errorf("check import work: %w", err)
 	}
 
-	old, err := rebuildReadLog(ctx, root)
+	old, err := rebuildReadLog(ctx, agent.dir)
 	if err != nil {
 		return fmt.Errorf("read import log: %w", err)
 	}
 
+	var rootLog []DatedMessage
+
+	if agent.dir != root {
+		rootLog, err = rebuildReadLog(ctx, root)
+		if err != nil {
+			return fmt.Errorf("read root import history: %w", err)
+		}
+	}
+
 	names, owner, err := whatsappImportNames(ctx, dir)
 	if err != nil {
 		return fmt.Errorf("read import names: %w", err)
 	}
 
-	plan, err := whatsappImportPlan(ctx, args, old, names, owner, time.Local)
+	plan, err := whatsappImportPlan(ctx, options.paths, old, rootLog, names, owner, time.Local)
 	if err != nil {
 		return fmt.Errorf("plan import: %w", err)
 	}
 
+	plan.agent = agent.config.Name
 	err = plan.Write(out)
 	if err != nil {
 		return fmt.Errorf("write import plan: %w", err)
 	}
 
-	if planOnly || len(plan.added) == 0 {
+	if options.planOnly || len(plan.added) == 0 {
 		return nil
 	}
 
 	// An existing agent is required so that a failed import never creates
 	// or changes the original agent's configuration.
-	_, err = os.Stat(filepath.Join(root, "config.json"))
+	_, err = os.Stat(filepath.Join(agent.dir, "config.json"))
 	if err != nil {
 		return fmt.Errorf("read agent config: %w", err)
 	}
 
-	agent, err := openAgent(root)
-	if err != nil {
-		return fmt.Errorf("open import agent: %w", err)
-	}
-
 	config, err := loadConfig(dir)
 	if err != nil {
 		return fmt.Errorf("load import configuration: %w", err)
@@ -142,7 +143,7 @@ func whatsappImport(ctx context.Context, args []string, out io.Writer) (err erro
 	}
 
 	client := newModelClient(address, auth, agent.CacheKey(), usagePath)
-	backup, err := rebuild(ctx, agent, plan.added, client, newCompactor(client, agent.Prompt, config.Compactor), out)
+	backup, err := rebuild(ctx, agent, plan.added, client, newCompactor(client, agent.Prompt, config.Compactor, agent.Tools(), agent.Images()), out)
 	if err != nil {
 		return fmt.Errorf("rebuild imported memory: %w", err)
 	}
@@ -159,6 +160,8 @@ func whatsappImport(ctx context.Context, args []string, out io.Writer) (err erro
 // log. Its estimate uses the measured cost of 1.3 calls per message, not a
 // promise about a particular export or compaction retry.
 type WhatsAppImportPlan struct {
+	agent      string
+	cutoffs    map[string]time.Time
 	added      []DatedMessage
 	chats      []string
 	cutoff     time.Time
@@ -187,13 +190,19 @@ func (p WhatsAppImportPlan) Write(out io.Writer) error {
 		owner = "unknown (sender names kept as written)"
 	}
 
-	_, err := fmt.Fprintf(out, "chats: %d\n", len(p.chats))
+	_, err := fmt.Fprintf(out, "agent: %s\nchats: %d\n", p.agent, len(p.chats))
 	if err != nil {
 		return fmt.Errorf("write import plan: %w", err)
 	}
 
 	for _, chat := range p.chats {
-		_, err = fmt.Fprintf(out, "  %s\n", chat)
+		line := "  " + chat
+		cutoff, found := p.cutoffs[chat]
+		if found {
+			line += " (cutoff: " + cutoff.Format(time.RFC3339) + ")"
+		}
+
+		_, err = fmt.Fprintln(out, line)
 		if err != nil {
 			return fmt.Errorf("write import plan: %w", err)
 		}
@@ -211,7 +220,7 @@ func (p WhatsAppImportPlan) Write(out io.Writer) error {
 		"owner: " + owner,
 		fmt.Sprintf("bytes: %d (new batch text)", p.bytes),
 		fmt.Sprintf("messages: %d (merged log)", p.messages),
-		fmt.Sprintf("estimated calls: %.0f (about 1.3 per message, limit %d)", estimate, buildCallLimit),
+		fmt.Sprintf("estimated calls: %.0f (about 1.3 per message, limit %d)", estimate, rebuildCallLimit),
 	}
 
 	_, err = io.WriteString(out, strings.Join(lines, "\n")+"\n")
@@ -222,6 +231,140 @@ func (p WhatsAppImportPlan) Write(out io.Writer) error {
 	return nil
 }
 
+// WhatsAppImportOptions selects the target, exports, and plan-only mode.
+type WhatsAppImportOptions struct {
+	planOnly bool
+	agent    string
+	paths    []string
+}
+
+// whatsappImportParse accepts GNU-style flags, including --agent=NAME and --
+// before paths that start with a dash.
+func whatsappImportParse(args []string) (WhatsAppImportOptions, error) {
+	var options WhatsAppImportOptions
+
+	flags := true
+	for i := 0; i < len(args); i++ {
+		arg := args[i]
+		switch {
+		case flags && arg == "--":
+			flags = false
+		case flags && arg == "-n":
+			options.planOnly = true
+		case flags && (arg == "--agent" || strings.HasPrefix(arg, "--agent=")):
+			_, name, found := strings.Cut(arg, "=")
+			if !found {
+				i++
+				if i == len(args) {
+					return options, errors.New("--agent needs a child name")
+				}
+
+				name = args[i]
+			}
+
+			if name == "" {
+				return options, errors.New("--agent needs a child name")
+			}
+
+			options.agent = name
+		case flags && strings.HasPrefix(arg, "-"):
+			return options, fmt.Errorf("unknown import flag %q", arg)
+		default:
+			options.paths = append(options.paths, arg)
+		}
+	}
+
+	if len(options.paths) == 0 {
+		return options, errors.New("usage: onechat import [-n] [--agent NAME] ZIP [ZIP...]")
+	}
+
+	return options, nil
+}
+
+// whatsappImportAgent selects the root or an existing direct child without
+// creating a directory or config. A root without a config can have a plan,
+// using the default name.
+func whatsappImportAgent(root, name string) (*Agent, error) {
+	_, err := os.Stat(filepath.Join(root, "config.json"))
+	if errors.Is(err, os.ErrNotExist) && name == "" {
+		var config AgentConfig
+
+		config, err = defaultAgentConfig()
+		if err != nil {
+			return nil, fmt.Errorf("read default import agent: %w", err)
+		}
+
+		return &Agent{
+			dir:    root,
+			config: config,
+		}, nil
+	}
+
+	if err != nil {
+		return nil, fmt.Errorf("read root agent config: %w", err)
+	}
+
+	agent, err := openAgent(root)
+	if err != nil {
+		return nil, fmt.Errorf("open root agent: %w", err)
+	}
+
+	if name == "" {
+		return agent, nil
+	}
+
+	if agent.config.Parent != "" {
+		return nil, errors.New("import root cannot have a parent")
+	}
+
+	children, err := agent.Children()
+	if err != nil {
+		return nil, fmt.Errorf("read import children: %w", err)
+	}
+
+	for _, child := range children {
+		if child.config.Name == name {
+			return child, nil
+		}
+	}
+
+	return nil, fmt.Errorf("child %q does not exist", name)
+}
+
+// whatsappImportIdle refuses unfinished subagents and pending deliveries to the
+// target. Rebuilding renumbers its log, so saved delivery IDs must not refer to
+// it until their turns finish. The caller holds the token lock.
+func whatsappImportIdle(root string, agent *Agent) error {
+	err := rebuildIdle(agent.dir)
+	if err != nil {
+		return fmt.Errorf("check target work: %w", err)
+	}
+
+	data, err := os.ReadFile(filepath.Join(root, "messages.json"))
+	if errors.Is(err, os.ErrNotExist) {
+		return nil
+	}
+
+	if err != nil {
+		return fmt.Errorf("read agent deliveries: %w", err)
+	}
+
+	var queue DaemonMessageFile
+
+	err = json.Unmarshal(data, &queue)
+	if err != nil {
+		return fmt.Errorf("decode agent deliveries: %w", err)
+	}
+
+	for _, message := range queue.Pending {
+		if message.Target == agent.config.Name {
+			return fmt.Errorf("import refused: agent %q has pending deliveries", agent.config.Name)
+		}
+	}
+
+	return nil
+}
+
 // whatsappImportNames reads only local session data. Group names that the
 // library keeps only in memory can also come from logged or pending batches.
 // Ambiguous names have no ID. A missing session keeps sender names unchanged.
@@ -319,64 +462,24 @@ func whatsappImportName(names map[string]WhatsAppChat, chat WhatsAppChat) {
 	names[chat.Name] = chat
 }
 
-// whatsappImportPlan filters exports at the first message of the earliest live
-// batch and skips chat-days already imported. Each remaining chat-day becomes a
-// work message dated at its last exported message. Overlapping chat-days in the
-// inputs are refused, not silently combined or duplicated.
-func whatsappImportPlan(ctx context.Context, paths []string, old []DatedMessage, names map[string]WhatsAppChat, owner string, location *time.Location) (WhatsAppImportPlan, error) {
-	plan := WhatsAppImportPlan{ownerKnown: owner != ""}
+// whatsappImportPlan filters each chat at its earliest live message across the
+// target and root logs. Only imported days in the target are skipped. Each
+// remaining chat-day becomes a work message dated at its last exported message.
+// Overlapping input chat-days are refused.
+func whatsappImportPlan(ctx context.Context, paths []string, old, rootLog []DatedMessage, names map[string]WhatsAppChat, owner string, location *time.Location) (WhatsAppImportPlan, error) {
+	plan := WhatsAppImportPlan{
+		ownerKnown: owner != "",
+		cutoffs:    make(map[string]time.Time),
+	}
+
 	names = maps.Clone(names)
 	if names == nil {
 		names = make(map[string]WhatsAppChat)
 	}
 
-	imported := make(map[string]bool)
-	for _, entry := range old {
-		if entry.Kind != "work" || !strings.HasPrefix(entry.Text, "[WhatsApp] ") {
-			continue
-		}
-
-		first, _, _ := strings.Cut(entry.Text, "\n")
-		marked := strings.HasSuffix(first, whatsappImportMark)
-		first = strings.TrimSuffix(first, whatsappImportMark)
-		match := whatsappImportHeader.FindStringSubmatch(first)
-		name, day := "", ""
-		if len(match) > 0 {
-			name, day = match[1], match[3]
-			whatsappImportName(names, WhatsAppChat{
-				ID:    match[2],
-				Name:  name,
-				Group: strings.Contains(first, ", group "+match[2]+","),
-			})
-		} else if marked {
-			match = whatsappImportUnknown.FindStringSubmatch(first)
-			if len(match) > 0 {
-				name, day = match[1], match[2]
-			}
-		}
-
-		if name == "" {
-			return plan, errors.New("read WhatsApp history: invalid batch header")
-		}
-
-		if marked {
-			imported[name+"\n"+day] = true
-			continue
-		}
-
-		_, body, found := strings.Cut(entry.Text, "\n")
-		if !found || len(body) < len("15:04 ") {
-			return plan, errors.New("read WhatsApp history: missing first message time")
-		}
-
-		date, err := time.ParseInLocation("2006-01-02 15:04", day+" "+body[:5], location)
-		if err != nil {
-			return plan, fmt.Errorf("parse first live message time: %w", err)
-		}
-
-		if plan.cutoff.IsZero() || date.Before(plan.cutoff) {
-			plan.cutoff = date
-		}
+	imported, err := whatsappImportHistory(old, rootLog, names, plan.cutoffs, location)
+	if err != nil {
+		return plan, fmt.Errorf("read import history: %w", err)
 	}
 
 	seen := make(map[string]bool)
@@ -399,9 +502,14 @@ func whatsappImportPlan(ctx context.Context, paths []string, old []DatedMessage,
 			plan.chats = append(plan.chats, chat.Name)
 		}
 
+		cutoff := plan.cutoffs[chat.Name]
+		if !cutoff.IsZero() && (plan.cutoff.IsZero() || cutoff.Before(plan.cutoff)) {
+			plan.cutoff = cutoff
+		}
+
 		days := make(map[string][]WhatsAppMessage)
 		for _, message := range messages {
-			if !plan.cutoff.IsZero() && !message.Time.Before(plan.cutoff) {
+			if !cutoff.IsZero() && !message.Time.Before(cutoff) {
 				continue
 			}
 
@@ -451,8 +559,6 @@ func whatsappImportPlan(ctx context.Context, paths []string, old []DatedMessage,
 		}
 	}
 
-	var err error
-
 	plan.added, err = rebuildMerge(nil, plan.added)
 	if err != nil {
 		return plan, fmt.Errorf("order import batches: %w", err)
@@ -464,6 +570,66 @@ func whatsappImportPlan(ctx context.Context, paths []string, old []DatedMessage,
 	return plan, nil
 }
 
+// whatsappImportHistory reads batch names and live cutoffs from both logs.
+// Imported chat-days belong only to the target log, which is first.
+func whatsappImportHistory(target, root []DatedMessage, names map[string]WhatsAppChat, cutoffs map[string]time.Time, location *time.Location) (map[string]bool, error) {
+	imported := make(map[string]bool)
+	for source, entries := range [][]DatedMessage{target, root} {
+		for _, entry := range entries {
+			if entry.Kind != "work" || !strings.HasPrefix(entry.Text, "[WhatsApp] ") {
+				continue
+			}
+
+			first, _, _ := strings.Cut(entry.Text, "\n")
+			marked := strings.HasSuffix(first, whatsappImportMark)
+			first = strings.TrimSuffix(first, whatsappImportMark)
+			match := whatsappImportHeader.FindStringSubmatch(first)
+			name, day := "", ""
+			if len(match) > 0 {
+				name, day = match[1], match[3]
+				whatsappImportName(names, WhatsAppChat{
+					ID:    match[2],
+					Name:  name,
+					Group: strings.Contains(first, ", group "+match[2]+","),
+				})
+			} else if marked {
+				match = whatsappImportUnknown.FindStringSubmatch(first)
+				if len(match) > 0 {
+					name, day = match[1], match[2]
+				}
+			}
+
+			if name == "" {
+				return nil, errors.New("invalid WhatsApp batch header")
+			}
+
+			if marked {
+				if source == 0 {
+					imported[name+"\n"+day] = true
+				}
+
+				continue
+			}
+
+			_, body, found := strings.Cut(entry.Text, "\n")
+			if !found || len(body) < len("15:04 ") {
+				return nil, errors.New("missing first WhatsApp message time")
+			}
+
+			date, err := time.ParseInLocation("2006-01-02 15:04", day+" "+body[:5], location)
+			if err != nil {
+				return nil, fmt.Errorf("parse first live message time: %w", err)
+			}
+
+			if cutoffs[name].IsZero() || date.Before(cutoffs[name]) {
+				cutoffs[name] = date
+			}
+		}
+	}
+
+	return imported, nil
+}
+
 // whatsappImportRead reads the single _chat.txt in an export without extracting
 // archive paths. Non-start lines continue the preceding message, and all
 // invisible left-to-right marks are removed. Dates use location. Media and