compact: structured TaskState compaction with recent message tail

pkg/agent/session.go

@@ -12,20 +12,23 @@ import (
 )
 
 const (
-	compactionPrompt = `You are performing a CONTEXT CHECKPOINT COMPACTION. Create a handoff summary for another LLM that will resume the task.
+	compactionPrompt = `You are performing a CONTEXT CHECKPOINT COMPACTION. Produce a JSON object summarizing the current task state. Output ONLY valid JSON matching this schema:
 
-Include:
-- Current progress and key decisions made
-- Important context, constraints, or user preferences
-- What remains to be done (clear next steps)
-- Any critical data, examples, or references needed to continue
+{
+  "objective": "the user's overall goal",
+  "plan_step": "what step you are currently on",
+  "constraints": ["any constraints or user preferences"],
+  "last_action": "brief: last tool/action + outcome",
+  "open_questions": ["unresolved questions"],
+  "next_decision": "what to do next"
+}
 
-Be concise, structured, and focused on helping the next LLM seamlessly continue the work.`
+Be concise. Capture what another LLM needs to seamlessly continue.`
 
-	compactionSummaryPrefix = `Another language model started to solve this problem and produced a summary of its thinking process. You also have access to the state of the tools that were used by that language model. Use this to build on the work that has already been done and avoid duplicating work.
-
-Here is the summary produced by the other language model, use the information in this summary to assist with your own analysis:`
+	compactionSummaryPrefix = `Context was compacted. Here is the structured task state from the prior context, followed by the most recent messages for evidential context:`
 
+	// recentTailSize is the number of messages to preserve after compaction.
+	recentTailSize = 6
 )
 
 // PersistedSession is the on-disk format for a saved session.
@@ -275,18 +278,61 @@ func (s *Session) compact(ctx context.Context, b backend.Backend) (int, string,
 		return 0, "", fmt.Errorf("compact: empty summary")
 	}
 
-	// Rebuild: system messages + preserved user messages + summary.
+	// Parse structured task state from the model's JSON response.
+	var ts TaskState
+	raw := extractJSON(summaryText)
+	if err := json.Unmarshal([]byte(raw), &ts); err != nil {
+		// Fallback: treat as unstructured summary if JSON parse fails.
+		ts = TaskState{Objective: summaryText}
+	}
+	s.TaskState = &ts
+
+	// Keep the last N messages as evidential context.
+	tail := s.messages
+	if len(tail) > recentTailSize {
+		tail = tail[len(tail)-recentTailSize:]
+	}
+
 	beforeCount := len(s.messages)
-	s.messages = buildCompactedHistory(summaryText)
+	s.messages = buildCompactedHistory(ts, tail)
 	return beforeCount - len(s.messages), summaryText, nil
 }
 
-// buildCompactedHistory constructs the post-compaction message list.
-func buildCompactedHistory(summary string) []backend.Message {
-	return []backend.Message{{
+// buildCompactedHistory constructs the post-compaction message list:
+// a structured task state header followed by recent messages as evidence.
+func buildCompactedHistory(ts TaskState, recentMessages []backend.Message) []backend.Message {
+	stateJSON, _ := json.Marshal(ts)
+	header := backend.Message{
 		Role:    "user",
-		Content: compactionSummaryPrefix + "\n" + strings.TrimSpace(summary),
-	}}
+		Content: compactionSummaryPrefix + "\n```json\n" + string(stateJSON) + "\n```",
+	}
+	return append([]backend.Message{header}, recentMessages...)
+}
+
+// extractJSON finds the first JSON object in s (handling markdown fences).
+func extractJSON(s string) string {
+	// Strip markdown code fence if present.
+	if i := strings.Index(s, "```json"); i >= 0 {
+		s = s[i+7:]
+		if j := strings.Index(s, "```"); j >= 0 {
+			s = s[:j]
+		}
+	} else if i := strings.Index(s, "```"); i >= 0 {
+		s = s[i+3:]
+		if j := strings.Index(s, "```"); j >= 0 {
+			s = s[:j]
+		}
+	}
+	// Find first { ... }
+	start := strings.Index(s, "{")
+	if start < 0 {
+		return s
+	}
+	end := strings.LastIndex(s, "}")
+	if end < start {
+		return s
+	}
+	return s[start : end+1]
 }
 
 // flattenToolMessages converts tool call/result sequences into plain text