CON-128: hybrid Post Assistant — cheap Haiku planner + Sonnet editPost writer

src/config/config.go

@@ -94,6 +94,23 @@ type Config struct {
 	// flow with explanation + full post content + tool inputs combined).
 	MaxOutputTokens int64 `envconfig:"MAX_OUTPUT_TOKENS"     default:"64000"`
 
+	// PostAssistantPlanner (CON-128) enables the hybrid model split for the
+	// Post Assistant: the orchestration/routing loop runs on the cheap
+	// PlanningModelID (Haiku) while the actual copywriting is delegated to a
+	// Sonnet (ModelID) editPost write-tool. Default on. Set to false to force
+	// the whole assistant back onto the proven single-Sonnet path (loop on
+	// ModelID, no editPost tool, inline content) — the instant rollback lever
+	// if Haiku routing regresses. Model ids stay tunable via MODEL_ID /
+	// PLANNING_MODEL_ID regardless.
+	PostAssistantPlanner bool `envconfig:"POST_ASSISTANT_PLANNER" default:"true"`
+
+	// PostAssistantPlannerMaxOutputTokens caps the Haiku planner turn's output
+	// (CON-128). The planner only emits a short envelope (explanation + action
+	// + saveVersion + versionNote) plus tool inputs, so a small cap is plenty;
+	// the full post is produced by the writer sub-call under MaxOutputTokens.
+	// 0 falls back to a sensible default (8192).
+	PostAssistantPlannerMaxOutputTokens int64 `envconfig:"POST_ASSISTANT_PLANNER_MAX_OUTPUT_TOKENS" default:"8192"`
+
 	// Content-plan batching. The flow generates posts in K-sized batches in
 	// parallel; the defaults are sized so a 64K-output Sonnet call comfortably
 	// returns 30 posts with headroom, and so an account with default tier

src/genkit/flows/post_assistant/edit_tool_test.go

@@ -0,0 +1,96 @@
+package post_assistant
+
+import (
+	"context"
+	"strings"
+	"testing"
+)
+
+// The editPost tool must reject an empty instruction before it ever spins up
+// the writer sub-call, and must not record an edit result.
+func TestToolEditPost_RequiresInstruction(t *testing.T) {
+	st := &requestState{postID: "p1"}
+	ctx := withRequestState(context.Background(), st)
+	if _, err := toolEditPost(ctx, EditPostInput{Instruction: "   "}); err == nil {
+		t.Fatal("expected an error for an empty instruction")
+	}
+	if st.editResult != nil {
+		t.Fatal("editResult must stay nil when no writer ran")
+	}
+}
+
+// runWriter is a no-op in the legacy state (no genkit instance / provider /
+// writer system prompt) — the guard keeps a stray call from panicking and
+// clearly reports the writer is unavailable.
+func TestRunWriter_Unavailable(t *testing.T) {
+	st := &requestState{postID: "p1"} // g / provider / writerSystem all zero
+	if _, err := runWriter(context.Background(), st, "shorten it", false); err == nil {
+		t.Fatal("expected an error when the writer is unavailable")
+	}
+}
+
+// A well-formed editPost call whose writer is unavailable surfaces a wrapped
+// write-content error and leaves editResult unset, so the runner never
+// finalises a bogus "edited" turn.
+func TestToolEditPost_WriterUnavailable(t *testing.T) {
+	st := &requestState{postID: "p1"} // writerSystem empty → runWriter errors
+	ctx := withRequestState(context.Background(), st)
+	_, err := toolEditPost(ctx, EditPostInput{Instruction: "make it punchier"})
+	if err == nil {
+		t.Fatal("expected an error when content writing is unavailable")
+	}
+	if !strings.Contains(err.Error(), "write content") {
+		t.Fatalf("expected a wrapped write-content error, got: %v", err)
+	}
+	if st.editResult != nil {
+		t.Fatal("editResult must stay nil on writer failure")
+	}
+}
+
+// The writer must receive the full retrieved excerpts as source material —
+// alongside the unchanged, verbatim instruction — so an asset-grounded edit
+// grounds on the retrieved text rather than the short preview (CON-128).
+func TestComposeWriterInstruction_IncludesRetrievedExcerpts(t *testing.T) {
+	instruction := "Add a section on goroutine scheduling from the whitepaper."
+	excerpts := []retrievedExcerpt{
+		{AssetID: "asset1", ChunkID: "c1", Content: "Goroutines are multiplexed onto OS threads by the Go runtime scheduler."},
+	}
+
+	out := composeWriterInstruction(instruction, excerpts)
+
+	if !strings.HasPrefix(out, instruction) {
+		t.Fatalf("the verbatim instruction must lead the writer prompt; got: %q", out)
+	}
+	if !strings.Contains(out, "multiplexed onto OS threads") {
+		t.Fatalf("the retrieved excerpt content must reach the writer as source material; got: %q", out)
+	}
+	if !strings.Contains(out, "Source material") {
+		t.Fatalf("excerpts should be labelled as source material; got: %q", out)
+	}
+
+	// No retrieval this turn → the instruction is passed through untouched.
+	if got := composeWriterInstruction(instruction, nil); got != instruction {
+		t.Fatalf("with no excerpts the instruction must be unchanged; got: %q", got)
+	}
+}
+
+// Asset-retrieval tools may return the same chunk more than once across a turn;
+// captureExcerpts must record each chunk once so the writer prompt isn't padded
+// with duplicates.
+func TestCaptureExcerpts_DedupesByChunkID(t *testing.T) {
+	st := &requestState{}
+	out := &ChunksOutput{Chunks: []ChunkContent{
+		{ID: "c1", Content: "alpha"},
+		{ID: "c2", Content: "beta"},
+	}}
+
+	st.captureExcerpts("a1", out)
+	st.captureExcerpts("a1", out) // same chunks retrieved again
+
+	if len(st.retrieved) != 2 {
+		t.Fatalf("expected 2 deduped excerpts, got %d", len(st.retrieved))
+	}
+	if st.retrieved[0].Content != "alpha" || st.retrieved[1].Content != "beta" {
+		t.Fatalf("captured excerpts lost their content: %+v", st.retrieved)
+	}
+}

src/genkit/flows/post_assistant/flow.go

@@ -5,6 +5,7 @@ import (
 	"embed"
 	"fmt"
 	"log/slog"
+	"strings"
 	"text/template"
 
 	"github.com/firebase/genkit/go/core"
@@ -38,9 +39,27 @@ func InitPostAssistant(g *genkit.Genkit, cfg PostAssistantFlowConfig, repos Post
 		return fmt.Errorf("parse post_assistant.tmpl: %w", err)
 	}
 	systemTmpl := tmpl.Lookup("system")
+	plannerTmpl := tmpl.Lookup("planner")
+	writerTmpl := tmpl.Lookup("writer")
 	contextTmpl := tmpl.Lookup("context")
-	if systemTmpl == nil || contextTmpl == nil {
-		return fmt.Errorf("post_assistant.tmpl must define both {{define \"system\"}} and {{define \"context\"}} blocks")
+	if systemTmpl == nil || plannerTmpl == nil || writerTmpl == nil || contextTmpl == nil {
+		return fmt.Errorf("post_assistant.tmpl must define \"system\", \"planner\", \"writer\", and \"context\" blocks")
+	}
+
+	// CON-128: in the hybrid path the orchestration loop runs on the planner
+	// system prompt (routing + the editPost tool) and delegates copywriting to
+	// the Sonnet writer; the legacy path keeps the single system prompt that
+	// writes content inline. Resolve which system prompt the loop uses, and
+	// render the (static) writer instructions once for the writer sub-call.
+	activeSystemTmpl := systemTmpl
+	var writerInstructions string
+	if cfg.PlannerEnabled {
+		activeSystemTmpl = plannerTmpl
+		wi, err := renderTemplate(writerTmpl, contextTemplateData{})
+		if err != nil {
+			return fmt.Errorf("render writer instructions: %w", err)
+		}
+		writerInstructions = strings.TrimSpace(wi)
 	}
 
 	tools := defineTools(g)
@@ -53,12 +72,12 @@ func InitPostAssistant(g *genkit.Genkit, cfg PostAssistantFlowConfig, repos Post
 
 	PostAssistantFlow = genkit.DefineFlow(g, "postAssistant",
 		func(ctx context.Context, req PostAssistantRequest) (*PostAssistantResponse, error) {
-			return runPostAssistant(ctx, g, req, cfg, repos, systemTmpl, contextTmpl, tools, nil)
+			return runPostAssistant(ctx, g, req, cfg, repos, activeSystemTmpl, contextTmpl, writerInstructions, tools, nil)
 		},
 	)
 
 	postAssistantRunner = func(ctx context.Context, req PostAssistantRequest, onEvent OnEventFunc) (*PostAssistantResponse, error) {
-		return runPostAssistant(ctx, g, req, cfg, repos, systemTmpl, contextTmpl, tools, onEvent)
+		return runPostAssistant(ctx, g, req, cfg, repos, activeSystemTmpl, contextTmpl, writerInstructions, tools, onEvent)
 	}
 
 	return nil

src/genkit/flows/post_assistant/prewarm.go

@@ -28,15 +28,26 @@ func prewarmToolCache(g *genkit.Genkit, cfg PostAssistantFlowConfig, t *toolSet)
 	ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
 	defer cancel()
 
+	// CON-128: warm the grammar the real loop uses — in the hybrid path that's
+	// the planner tool set (incl. editPost) on the planning model; in the legacy
+	// path it's the base tool set on the generation model. The writer sub-call
+	// carries no tools, so there is no separate grammar to warm for it.
+	role := llm.RoleGeneration
+	tools := []ai.ToolRef{
+		t.listAssets, t.getAssetChunks, t.searchAssetChunks, t.getCurrentContent,
+		t.clonePost, t.restoreVersion, t.schedulePost, t.createNote,
+	}
+	if cfg.PlannerEnabled {
+		role = llm.RolePlanning
+		tools = append(tools, t.editPost)
+	}
+
 	start := time.Now()
 	_, err := genkit.Generate(ctx, g,
-		ai.WithModelName(cfg.Provider.Ref(llm.RoleGeneration)),
+		ai.WithModelName(cfg.Provider.Ref(role)),
 		ai.WithSystem("warmup"),
 		ai.WithPrompt("warmup"),
-		ai.WithTools(
-			t.listAssets, t.getAssetChunks, t.searchAssetChunks, t.getCurrentContent,
-			t.clonePost, t.restoreVersion, t.schedulePost, t.createNote,
-		),
+		ai.WithTools(tools...),
 		ai.WithMaxTurns(1),
 		cfg.Provider.CallConfig(1), // max_tokens: 1 — grammar compiles during prep
 	)