plugin/guardrail/approval: synchronize tool-call context before review

agent/invocation.go

@@ -57,6 +57,9 @@ const (
 	flusherStateKey = "__flush_session__"
 	// barrierStateKey is the invocation state key used by internal barrier flag.
 	barrierStateKey = "__graph_barrier__"
+	// toolCallPersistenceStateKey carries the barrier requirement for tool
+	// callbacks that read the current assistant response from the session.
+	toolCallPersistenceStateKey = "__tool_call_persistence__"
 	// appenderStateKey is the invocation state key used by internal appender
 	// attachment (see internal/state/appender).
 	appenderStateKey = "__append_event__"
@@ -1752,6 +1755,7 @@ func isCloneStateKey(key string) bool {
 	switch key {
 	case flusherStateKey,
 		barrierStateKey,
+		toolCallPersistenceStateKey,
 		appenderStateKey,
 		liveSessionStateKey,
 		streamHubStateKey,

examples/guardrail/approval/README.md

@@ -13,6 +13,24 @@ This example shows how to wire the top-level `guardrail` plugin with the built-i
   - `hostexec_kill_session` is denied by policy
 - Observe approval logs for approved and denied review decisions
 
+With a reviewer configured, the approval plugin makes runner-managed LLMAgents
+wait for the completed assistant/tool-call event's persistence acknowledgement
+before tool callbacks run. This includes serial and parallel tools and candidate
+attempts, so the reviewer can read the current assistant message and all tool
+calls in that response from the session. Existing transcript filters and size
+limits still apply. If synchronization fails, tool processing stops.
+
+Partial responses remain streaming. Standalone agents without a completion
+consumer and runs without a configured reviewer retain their existing behavior.
+Approval still runs in `BeforeTool`; this does not change callback ordering or
+the `Reviewer` interface.
+
+For native Graph LLM nodes, enable
+`agent.WithGraphEmitFinalModelResponses(true)` to make completed model responses
+available in the session transcript. Graph's existing node barrier orders those
+events before a subsequent tools node; its default response emission policy is
+unchanged.
+
 ## Prerequisites
 
 - Go 1.21 or later

internal/flow/llmflow/llmflow.go

@@ -37,6 +37,7 @@ import (
 	"trpc.group/trpc-go/trpc-agent-go/internal/modelcontext"
 	imodelrequest "trpc.group/trpc-go/trpc-agent-go/internal/modelrequest"
 	"trpc.group/trpc-go/trpc-agent-go/internal/responseusage"
+	"trpc.group/trpc-go/trpc-agent-go/internal/state/barrier"
 	"trpc.group/trpc-go/trpc-agent-go/internal/state/steer"
 	"trpc.group/trpc-go/trpc-agent-go/internal/state/summaryfork"
 	"trpc.group/trpc-go/trpc-agent-go/internal/state/summaryinject"
@@ -1002,11 +1003,16 @@ func (p *streamingResponseProcessor) process(
 		}
 		return true
 	}
-	llmResponseEvent := p.emitLLMResponse(
+	llmResponseEvent, err := p.emitLLMResponse(
 		eventInvocation,
 		response,
 		traceDetails,
 	)
+	if err != nil {
+		*p.err = err
+		responseErr = err
+		return false
+	}
 	p.lastEvent = llmResponseEvent
 	if p.tracker != nil {
 		p.tracker.SetLastEvent(p.lastEvent)
@@ -1191,13 +1197,22 @@ func (p *streamingResponseProcessor) emitLLMResponse(
 	eventInvocation *agent.Invocation,
 	response *model.Response,
 	traceDetails bool,
-) *event.Event {
+) (*event.Event, error) {
 	llmResponseEvent := p.flow.createLLMResponseEvent(
 		eventInvocation,
 		p.currentInvocation,
 		response,
 		p.llmRequest,
 	)
+	// Agent wrappers forward events asynchronously. BeforeTool reviewers need
+	// this response in the session before callbacks run or parallel tools clone
+	// the session. Wait for this event's acknowledgement, not a channel drain.
+	waitForPersistence := barrier.Enabled(eventInvocation) &&
+		barrier.ToolCallPersistenceRequired(eventInvocation) &&
+		!response.IsPartial && response.IsToolCallResponse()
+	if waitForPersistence {
+		llmResponseEvent.RequiresCompletion = true
+	}
 	emitCtx := p.ctx
 	var emitSpan oteltrace.Span
 	emitStarted := false
@@ -1209,9 +1224,21 @@ func (p *streamingResponseProcessor) emitLLMResponse(
 			latencyResponseAttrs(response)...,
 		)
 	}
-	agent.EmitEvent(emitCtx, eventInvocation, p.eventChan, llmResponseEvent)
-	finishLatencySpan(emitSpan, emitStarted, nil)
-	return llmResponseEvent
+	err := agent.EmitEvent(emitCtx, eventInvocation, p.eventChan, llmResponseEvent)
+	if !waitForPersistence {
+		// Preserve emission behavior for runs without this synchronization need.
+		finishLatencySpan(emitSpan, emitStarted, nil)
+		return llmResponseEvent, nil
+	}
+	if err == nil {
+		err = eventInvocation.AddNoticeChannelAndWait(p.ctx,
+			agent.GetAppendEventNoticeKey(llmResponseEvent.ID), flowEventWaitTimeout(p.ctx))
+	}
+	finishLatencySpan(emitSpan, emitStarted, err)
+	if err != nil {
+		return nil, fmt.Errorf("wait for tool-call event persistence: %w", err)
+	}
+	return llmResponseEvent, nil
 }
 
 func (p *streamingResponseProcessor) traceChat(

internal/flow/llmflow/tool_call_persistence_test.go

@@ -0,0 +1,107 @@
+//
+// Tencent is pleased to support the open source community by making trpc-agent-go available.
+//
+// Copyright (C) 2026 Tencent.  All rights reserved.
+//
+// trpc-agent-go is licensed under the Apache License Version 2.0.
+//
+
+package llmflow
+
+import (
+	"context"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	oteltrace "go.opentelemetry.io/otel/trace"
+	"trpc.group/trpc-go/trpc-agent-go/agent"
+	"trpc.group/trpc-go/trpc-agent-go/event"
+	"trpc.group/trpc-go/trpc-agent-go/internal/flow"
+	"trpc.group/trpc-go/trpc-agent-go/internal/state/barrier"
+	"trpc.group/trpc-go/trpc-agent-go/model"
+)
+
+type persistenceObserver struct {
+	called                atomic.Bool
+	acknowledged          *atomic.Bool
+	beforeAcknowledgement atomic.Bool
+}
+
+func (p *persistenceObserver) ProcessResponse(context.Context, *agent.Invocation, *model.Request, *model.Response, chan<- *event.Event) {
+	p.called.Store(true)
+	p.beforeAcknowledgement.Store(!p.acknowledged.Load())
+}
+
+func TestToolCallPersistenceBeforeResponseProcessors(t *testing.T) {
+	for _, tc := range []struct {
+		name                                          string
+		required, consumer, partial, toolCall, cancel bool
+		wantWait                                      bool
+	}{
+		{name: "acknowledged", required: true, consumer: true, toolCall: true, wantWait: true},
+		{name: "canceled before acknowledgement", required: true, consumer: true, toolCall: true, cancel: true, wantWait: true},
+		{name: "no requirement", consumer: true, toolCall: true},
+		{name: "standalone flow", required: true, toolCall: true},
+		{name: "partial tool call", required: true, consumer: true, partial: true, toolCall: true},
+		{name: "text response", required: true, consumer: true},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+			defer cancel()
+			inv := agent.NewInvocation()
+			if tc.required {
+				barrier.RequireToolCallPersistence(inv)
+			}
+			if tc.consumer {
+				barrier.Enable(inv)
+			}
+			var acknowledged atomic.Bool
+			post := &persistenceObserver{acknowledged: &acknowledged}
+			f := New(nil, []flow.ResponseProcessor{post}, Options{})
+			msg := model.NewAssistantMessage("Inspect the file.")
+			if tc.toolCall {
+				msg.ToolCalls = []model.ToolCall{{ID: "call", Function: model.FunctionDefinitionParam{Name: "read", Arguments: []byte(`{}`)}}}
+			}
+			rsp := &model.Response{IsPartial: tc.partial, Choices: []model.Choice{{Message: msg}}}
+			out := make(chan *event.Event, 1)
+			var processErr error
+			p := newStreamingResponseProcessor(f, ctx, inv, nil, &model.Request{}, out,
+				oteltrace.SpanFromContext(ctx), false, &processErr)
+			done := make(chan bool, 1)
+			go func() { done <- p.process(rsp) }()
+			var evt *event.Event
+			select {
+			case evt = <-out:
+			case <-ctx.Done():
+				t.Fatal(ctx.Err())
+			}
+			require.Equal(t, tc.wantWait, evt.RequiresCompletion)
+			if tc.wantWait {
+				require.False(t, post.called.Load())
+				if tc.cancel {
+					cancel()
+				} else {
+					acknowledged.Store(true)
+					require.NoError(t, inv.NotifyCompletion(ctx, agent.GetAppendEventNoticeKey(evt.ID)))
+				}
+			}
+			select {
+			case ok := <-done:
+				if tc.cancel {
+					require.False(t, ok)
+					require.ErrorIs(t, processErr, context.Canceled)
+					require.False(t, post.called.Load(), "review and tool processors must not run after synchronization fails")
+				} else {
+					require.True(t, ok)
+					require.NoError(t, processErr)
+					require.True(t, post.called.Load())
+					require.Equal(t, !tc.wantWait, post.beforeAcknowledgement.Load())
+				}
+			case <-time.After(5 * time.Second):
+				t.Fatal("response processing did not finish")
+			}
+		})
+	}
+}

internal/state/barrier/barrier.go

@@ -7,7 +7,7 @@
 //
 //
 
-// Package barrier provides a mechanism to indicate graph barrier requirements.
+// Package barrier records invocation event synchronization requirements.
 package barrier
 
 import (
@@ -17,6 +17,8 @@ import (
 // StateKeyBarrier is the invocation state key used by internal barrier flag.
 const StateKeyBarrier = "__graph_barrier__"
 
+const toolCallPersistenceKey = "__tool_call_persistence__"
+
 // Enable enables the graph barrier for the invocation.
 func Enable(inv *agent.Invocation) {
 	if inv == nil {
@@ -33,3 +35,20 @@ func Enabled(inv *agent.Invocation) bool {
 	}
 	return enabled
 }
+
+// RequireToolCallPersistence requests that completed LLM tool-call events be
+// acknowledged before tool processing. The requirement is inherited by cloned
+// invocations and only takes effect when Enabled reports a completion consumer.
+// A nil invocation is ignored.
+func RequireToolCallPersistence(inv *agent.Invocation) {
+	if inv != nil {
+		inv.SetState(toolCallPersistenceKey, true)
+	}
+}
+
+// ToolCallPersistenceRequired reports whether the invocation requests tool-call
+// persistence synchronization. Nil and unmarked invocations return false.
+func ToolCallPersistenceRequired(inv *agent.Invocation) bool {
+	required, _ := agent.GetStateValue[bool](inv, toolCallPersistenceKey)
+	return required
+}

internal/state/barrier/barrier_test.go

@@ -31,3 +31,14 @@ func TestEnableNilInvocation(t *testing.T) {
 	Enable(nil)
 	require.False(t, Enabled(nil))
 }
+
+func TestToolCallPersistenceRequirement(t *testing.T) {
+	RequireToolCallPersistence(nil)
+	require.False(t, ToolCallPersistenceRequired(nil))
+	inv := agent.NewInvocation()
+	require.False(t, ToolCallPersistenceRequired(inv))
+	RequireToolCallPersistence(inv)
+	require.True(t, ToolCallPersistenceRequired(inv))
+	require.True(t, ToolCallPersistenceRequired(inv.Clone()))
+	require.False(t, Enabled(inv), "requesting synchronization must not advertise a completion consumer")
+}

plugin/guardrail/approval/approval.go

@@ -14,6 +14,8 @@ import (
 	"fmt"
 	"strings"
 
+	"trpc.group/trpc-go/trpc-agent-go/agent"
+	"trpc.group/trpc-go/trpc-agent-go/internal/state/barrier"
 	"trpc.group/trpc-go/trpc-agent-go/log"
 	"trpc.group/trpc-go/trpc-agent-go/model"
 	"trpc.group/trpc-go/trpc-agent-go/plugin"
@@ -66,6 +68,14 @@ func (p *Plugin) Register(r *plugin.Registry) {
 	if p == nil || r == nil {
 		return
 	}
+	if p.reviewer != nil {
+		r.BeforeAgent(func(_ context.Context, args *agent.BeforeAgentArgs) (*agent.BeforeAgentResult, error) {
+			if args != nil {
+				barrier.RequireToolCallPersistence(args.Invocation)
+			}
+			return nil, nil
+		})
+	}
 	r.BeforeTool(p.beforeTool())
 }
 

plugin/guardrail/approval/approval_test.go

@@ -18,6 +18,7 @@ import (
 	"github.com/stretchr/testify/require"
 	"trpc.group/trpc-go/trpc-agent-go/agent"
 	"trpc.group/trpc-go/trpc-agent-go/event"
+	"trpc.group/trpc-go/trpc-agent-go/internal/state/barrier"
 	approvallog "trpc.group/trpc-go/trpc-agent-go/log"
 	"trpc.group/trpc-go/trpc-agent-go/model"
 	"trpc.group/trpc-go/trpc-agent-go/plugin"
@@ -101,6 +102,24 @@ func TestRegister_IgnoresNilReceiverAndNilRegistry(t *testing.T) {
 	})
 }
 
+func TestRegister_RequestsToolCallPersistenceOnlyWithReviewer(t *testing.T) {
+	withoutReviewer, err := New(WithDefaultToolPolicy(ToolPolicySkipApproval))
+	require.NoError(t, err)
+	require.Nil(t, plugin.MustNewManager(withoutReviewer).AgentCallbacks())
+	withReviewer, err := New(WithReviewer(&stubReviewer{}))
+	require.NoError(t, err)
+	callbacks := plugin.MustNewManager(withReviewer).AgentCallbacks()
+	require.NotNil(t, callbacks)
+	_, err = callbacks.RunBeforeAgent(context.Background(), nil)
+	require.NoError(t, err)
+	inv := agent.NewInvocation()
+	result, err := callbacks.RunBeforeAgent(context.Background(), &agent.BeforeAgentArgs{Invocation: inv})
+	require.NoError(t, err)
+	require.Nil(t, result)
+	require.True(t, barrier.ToolCallPersistenceRequired(inv))
+	require.False(t, barrier.Enabled(inv))
+}
+
 func TestWithToolPolicy_InitializesMap(t *testing.T) {
 	opts := &options{}
 	WithToolPolicy("shell", ToolPolicyDenied)(opts)