From b775e5df7a5dfab3c78a0f3e039fc10b231b8e6f Mon Sep 17 00:00:00 2001
From: Levi Neely <141506390+lneely@users.noreply.github.com>
Date: Sun, 17 May 2026 10:06:21 +0200
Subject: [PATCH] =?UTF-8?q?agent:=20add=20TaskState=20=E2=80=94=20persiste?=
 =?UTF-8?q?d=20structured=20overlay=20injected=20at=20each=20turn?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 pkg/agent/core.go    |  6 +++++
 pkg/agent/loop.go    |  9 +++++++
 pkg/agent/session.go | 58 ++++++++++++++++++++++++++++++++++++++------
 3 files changed, 65 insertions(+), 8 deletions(-)

diff --git a/pkg/agent/core.go b/pkg/agent/core.go
index 340c60d..7804d81 100644
--- a/pkg/agent/core.go
+++ b/pkg/agent/core.go
@@ -969,6 +969,12 @@ func (s *agent) executeTurn(ctx context.Context, input string, handler EventHand
 	s.cfg.IncrToolCallCount = func() int64 {
 		return s.toolCallCount.Add(1)
 	}
+	s.cfg.TaskState = func() *TaskState {
+		if s.session == nil {
+			return nil
+		}
+		return s.session.TaskState
+	}
 	s.cfg.SaveSession = func() { s.saveSession() }
 	s.cfg.AutoCompact = func(ctx context.Context) {
 		if ctx.Err() != nil || s.session == nil {
diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go
index 99b84f2..bfe360a 100644
--- a/pkg/agent/loop.go
+++ b/pkg/agent/loop.go
@@ -62,6 +62,8 @@ type agentConfig struct {
 	PlanFile string
 	// IncrToolCallCount increments and returns the session-wide tool call counter.
 	IncrToolCallCount func() int64
+	// TaskState returns the current task state, or nil if not set.
+	TaskState func() *TaskState
 }
 
 // readPlanFile returns the contents of the plan file, or empty string if
@@ -90,6 +92,13 @@ func run(ctx context.Context, cfg agentConfig, state state) error {
 		if cfg.preamble != "" {
 			history = append([]backend.Message{{Role: "system", Content: cfg.preamble}}, history...)
 		}
+		if cfg.TaskState != nil {
+			if ts := cfg.TaskState(); ts != nil {
+				if msg := ts.render(); msg != "" {
+					history = append([]backend.Message{history[0], {Role: "user", Content: msg}}, history[1:]...)
+				}
+			}
+		}
 
 		// Stream the assistant's response, retrying on rate limits, transient
 		// backend errors (5xx, network), and mid-stream drops.
diff --git a/pkg/agent/session.go b/pkg/agent/session.go
index 07f4a21..fb815a4 100644
--- a/pkg/agent/session.go
+++ b/pkg/agent/session.go
@@ -30,17 +30,58 @@ Here is the summary produced by the other language model, use the information in
 
 // PersistedSession is the on-disk format for a saved session.
 type PersistedSession struct {
-	ID       string            `json:"id"`
-	Agent    string            `json:"agent,omitempty"`
-	Messages []backend.Message `json:"messages"`
+	ID        string            `json:"id"`
+	Agent     string            `json:"agent,omitempty"`
+	Messages  []backend.Message `json:"messages"`
+	TaskState *TaskState        `json:"taskState,omitempty"`
+}
+
+// TaskState is a compact structured overlay that summarizes the agent's
+// current position in a task. Injected at the top of every turn so the
+// model doesn't rely on full history recall.
+type TaskState struct {
+	Objective     string   `json:"objective"`
+	PlanStep      string   `json:"plan_step"`
+	Constraints   []string `json:"constraints,omitempty"`
+	LastAction    string   `json:"last_action"`
+	OpenQuestions []string `json:"open_questions,omitempty"`
+	NextDecision  string   `json:"next_decision"`
+}
+
+// render produces a compact text representation for injection into the message history.
+func (ts *TaskState) render() string {
+	if ts.Objective == "" {
+		return ""
+	}
+	var sb strings.Builder
+	sb.WriteString("[task state]\n")
+	sb.WriteString("objective: " + ts.Objective + "\n")
+	if ts.PlanStep != "" {
+		sb.WriteString("plan_step: " + ts.PlanStep + "\n")
+	}
+	if len(ts.Constraints) > 0 {
+		sb.WriteString("constraints: " + strings.Join(ts.Constraints, "; ") + "\n")
+	}
+	if ts.LastAction != "" {
+		sb.WriteString("last_action: " + ts.LastAction + "\n")
+	}
+	if len(ts.OpenQuestions) > 0 {
+		sb.WriteString("open_questions: " + strings.Join(ts.OpenQuestions, "; ") + "\n")
+	}
+	if ts.NextDecision != "" {
+		sb.WriteString("next_decision: " + ts.NextDecision + "\n")
+	}
+	sb.WriteString("[/task state]")
+	return sb.String()
 }
 
 // SaveTo writes the full message history to path as JSON.
 func (s *Session) saveTo(path, id, agentName string) error {
 	ps := PersistedSession{
-		ID:       id,
-		Agent:    agentName,
-		Messages: s.messages,
+		ID:        id,
+		Agent:     agentName,
+		Messages:  s.messages,
+		TaskState: s.TaskState,
 	}
 	data, err := json.Marshal(ps)
 	if err != nil {
@@ -63,8 +104,9 @@ func RestoreSession(messages []backend.Message) *Session {
 
 // Session is an ephemeral in-memory state backend.
 type Session struct {
-	goal     string
-	messages []backend.Message
+	goal      string
+	messages  []backend.Message
+	TaskState *TaskState
 	// Cumulative usage tracking.
 	TotalInputTokens         int
 	TotalCachedInputTokens   int