Skip to content

Commit 80e2dcd

Browse files
gustafvhgustafvh
committed
fix: yield function responses before tool confirmations
A consumer that pauses on receiving an adk_request_confirmation event (typically to await human approval) would never see the merged function-response event for tools that already executed in the same step. Yield completed responses first so they are persisted before the human-confirmation pause. Fixes #759
1 parent a586408 commit 80e2dcd

3 files changed

Lines changed: 100 additions & 43 deletions

File tree

‎internal/llminternal/base_flow.go‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -204,16 +204,19 @@ func (f *Flow) runOneStep(ctx agent.InvocationContext) iter.Seq2[*session.Event,
204204
}
205205

206206
toolConfirmationEvent := generateRequestConfirmationEvent(ctx, modelResponseEvent, ev)
207+
208+
// Yield function responses before confirmation requests so consumers that
209+
// pause for user approval still persist completed tool results.
210+
if !yield(ev, nil) {
211+
return
212+
}
213+
207214
if toolConfirmationEvent != nil {
208215
if !yield(toolConfirmationEvent, nil) {
209216
return
210217
}
211218
}
212219

213-
if !yield(ev, nil) {
214-
return
215-
}
216-
217220
// If the model response is structured, yield it as a final model response event.
218221
outputSchemaResponse, err := retrieveStructuredModelResponse(ev)
219222
if err != nil {

‎tool/functiontool/function_test.go‎

Lines changed: 72 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -584,6 +584,9 @@ func TestToolConfirmation(t *testing.T) {
584584
args: map[string]any{"Num": 1},
585585
want: []*genai.Content{
586586
genai.NewContentFromFunctionCall("test_tool", map[string]any{"Num": 1}, "model"),
587+
genai.NewContentFromFunctionResponse("test_tool", map[string]any{
588+
"error": errors.New("error tool \"test_tool\" requires confirmation, please approve or reject"),
589+
}, "user"),
587590
genai.NewContentFromFunctionCall(toolconfirmation.FunctionCallName, map[string]any{
588591
"originalFunctionCall": &genai.FunctionCall{
589592
Args: map[string]any{"Num": 1},
@@ -593,9 +596,6 @@ func TestToolConfirmation(t *testing.T) {
593596
Hint: "Please approve or reject the tool call test_tool() by responding with a FunctionResponse with an expected ToolConfirmation payload.",
594597
},
595598
}, "model"),
596-
genai.NewContentFromFunctionResponse("test_tool", map[string]any{
597-
"error": errors.New("error tool \"test_tool\" requires confirmation, please approve or reject"),
598-
}, "user"),
599599
},
600600
},
601601
{
@@ -608,6 +608,9 @@ func TestToolConfirmation(t *testing.T) {
608608
confirmFunctionResponse: &genai.FunctionResponse{Name: toolconfirmation.FunctionCallName, Response: map[string]any{"confirmed": true}},
609609
want: []*genai.Content{
610610
genai.NewContentFromFunctionCall("test_tool", map[string]any{"Num": 1}, "model"),
611+
genai.NewContentFromFunctionResponse("test_tool", map[string]any{
612+
"error": errors.New("error tool \"test_tool\" requires confirmation, please approve or reject"),
613+
}, "user"),
611614
genai.NewContentFromFunctionCall(toolconfirmation.FunctionCallName, map[string]any{
612615
"originalFunctionCall": &genai.FunctionCall{
613616
Args: map[string]any{"Num": 1},
@@ -617,9 +620,6 @@ func TestToolConfirmation(t *testing.T) {
617620
Hint: "Please approve or reject the tool call test_tool() by responding with a FunctionResponse with an expected ToolConfirmation payload.",
618621
},
619622
}, "model"),
620-
genai.NewContentFromFunctionResponse("test_tool", map[string]any{
621-
"error": errors.New("error tool \"test_tool\" requires confirmation, please approve or reject"),
622-
}, "user"),
623623
genai.NewContentFromFunctionResponse("test_tool", map[string]any{"result": "ok"}, "user"),
624624
},
625625
},
@@ -633,6 +633,9 @@ func TestToolConfirmation(t *testing.T) {
633633
confirmFunctionResponse: &genai.FunctionResponse{Name: toolconfirmation.FunctionCallName, Response: map[string]any{"confirmed": false}},
634634
want: []*genai.Content{
635635
genai.NewContentFromFunctionCall("test_tool", map[string]any{"Num": 1}, "model"),
636+
genai.NewContentFromFunctionResponse("test_tool", map[string]any{
637+
"error": errors.New("error tool \"test_tool\" requires confirmation, please approve or reject"),
638+
}, "user"),
636639
genai.NewContentFromFunctionCall(toolconfirmation.FunctionCallName, map[string]any{
637640
"originalFunctionCall": &genai.FunctionCall{
638641
Args: map[string]any{"Num": 1},
@@ -642,9 +645,6 @@ func TestToolConfirmation(t *testing.T) {
642645
Hint: "Please approve or reject the tool call test_tool() by responding with a FunctionResponse with an expected ToolConfirmation payload.",
643646
},
644647
}, "model"),
645-
genai.NewContentFromFunctionResponse("test_tool", map[string]any{
646-
"error": errors.New("error tool \"test_tool\" requires confirmation, please approve or reject"),
647-
}, "user"),
648648
genai.NewContentFromFunctionResponse("test_tool", map[string]any{
649649
"error": errors.New("error tool \"test_tool\" call is rejected"),
650650
}, "user"),
@@ -675,6 +675,9 @@ func TestToolConfirmation(t *testing.T) {
675675
args: map[string]any{"Num": 4},
676676
want: []*genai.Content{
677677
genai.NewContentFromFunctionCall("test_tool", map[string]any{"Num": 4}, "model"),
678+
genai.NewContentFromFunctionResponse("test_tool", map[string]any{
679+
"error": errors.New("error tool \"test_tool\" requires confirmation, please approve or reject"),
680+
}, "user"),
678681
genai.NewContentFromFunctionCall(toolconfirmation.FunctionCallName, map[string]any{
679682
"originalFunctionCall": &genai.FunctionCall{
680683
Args: map[string]any{"Num": 4},
@@ -684,9 +687,6 @@ func TestToolConfirmation(t *testing.T) {
684687
Hint: "Please approve or reject the tool call test_tool() by responding with a FunctionResponse with an expected ToolConfirmation payload.",
685688
},
686689
}, "model"),
687-
genai.NewContentFromFunctionResponse("test_tool", map[string]any{
688-
"error": errors.New("error tool \"test_tool\" requires confirmation, please approve or reject"),
689-
}, "user"),
690690
},
691691
},
692692
{
@@ -701,6 +701,9 @@ func TestToolConfirmation(t *testing.T) {
701701
confirmFunctionResponse: &genai.FunctionResponse{Name: toolconfirmation.FunctionCallName, Response: map[string]any{"confirmed": true}},
702702
want: []*genai.Content{
703703
genai.NewContentFromFunctionCall("test_tool", map[string]any{"Num": 4}, "model"),
704+
genai.NewContentFromFunctionResponse("test_tool", map[string]any{
705+
"error": errors.New("error tool \"test_tool\" requires confirmation, please approve or reject"),
706+
}, "user"),
704707
genai.NewContentFromFunctionCall(toolconfirmation.FunctionCallName, map[string]any{
705708
"originalFunctionCall": &genai.FunctionCall{
706709
Args: map[string]any{"Num": 4},
@@ -710,9 +713,6 @@ func TestToolConfirmation(t *testing.T) {
710713
Hint: "Please approve or reject the tool call test_tool() by responding with a FunctionResponse with an expected ToolConfirmation payload.",
711714
},
712715
}, "model"),
713-
genai.NewContentFromFunctionResponse("test_tool", map[string]any{
714-
"error": errors.New("error tool \"test_tool\" requires confirmation, please approve or reject"),
715-
}, "user"),
716716
genai.NewContentFromFunctionResponse("test_tool", map[string]any{"result": "ok"}, "user"),
717717
},
718718
},
@@ -728,6 +728,9 @@ func TestToolConfirmation(t *testing.T) {
728728
confirmFunctionResponse: &genai.FunctionResponse{Name: toolconfirmation.FunctionCallName, Response: map[string]any{"confirmed": false}},
729729
want: []*genai.Content{
730730
genai.NewContentFromFunctionCall("test_tool", map[string]any{"Num": 4}, "model"),
731+
genai.NewContentFromFunctionResponse("test_tool", map[string]any{
732+
"error": errors.New("error tool \"test_tool\" requires confirmation, please approve or reject"),
733+
}, "user"),
731734
genai.NewContentFromFunctionCall(toolconfirmation.FunctionCallName, map[string]any{
732735
"originalFunctionCall": &genai.FunctionCall{
733736
Args: map[string]any{"Num": 4},
@@ -737,9 +740,6 @@ func TestToolConfirmation(t *testing.T) {
737740
Hint: "Please approve or reject the tool call test_tool() by responding with a FunctionResponse with an expected ToolConfirmation payload.",
738741
},
739742
}, "model"),
740-
genai.NewContentFromFunctionResponse("test_tool", map[string]any{
741-
"error": errors.New("error tool \"test_tool\" requires confirmation, please approve or reject"),
742-
}, "user"),
743743
genai.NewContentFromFunctionResponse("test_tool", map[string]any{
744744
"error": errors.New("error tool \"test_tool\" call is rejected"),
745745
}, "user"),
@@ -865,6 +865,60 @@ func TestToolConfirmation(t *testing.T) {
865865
}
866866
}
867867

868+
func TestToolConfirmationYieldsFunctionResponseBeforeConfirmation(t *testing.T) {
869+
mockModel := &testutil.MockModel{
870+
Responses: []*genai.Content{
871+
genai.NewContentFromFunctionCall("test_tool", map[string]any{"Num": 1}, genai.RoleModel),
872+
},
873+
}
874+
875+
myTool, err := functiontool.New(functiontool.Config{
876+
Name: "test_tool",
877+
RequireConfirmation: true,
878+
}, okFunc)
879+
if err != nil {
880+
t.Fatalf("Failed to create tool: %v", err)
881+
}
882+
883+
a, err := llmagent.New(llmagent.Config{
884+
Name: "simple agent",
885+
Model: mockModel,
886+
Tools: []tool.Tool{myTool},
887+
})
888+
if err != nil {
889+
t.Fatalf("failed to create llm agent: %v", err)
890+
}
891+
892+
runner := testutil.NewTestAgentRunner(t, a)
893+
sawFunctionResponse := false
894+
sawConfirmation := false
895+
896+
for got, err := range runner.Run(t, "id", "message") {
897+
// The mock model emits a sentinel error once exhausted; treat any
898+
// error as end-of-stream and rely on the post-loop assertions.
899+
if err != nil {
900+
break
901+
}
902+
903+
for _, part := range got.Content.Parts {
904+
if part.FunctionResponse != nil && part.FunctionResponse.Name == "test_tool" {
905+
sawFunctionResponse = true
906+
}
907+
if part.FunctionCall != nil && part.FunctionCall.Name == toolconfirmation.FunctionCallName {
908+
sawConfirmation = true
909+
if !sawFunctionResponse {
910+
t.Fatal("confirmation was yielded before the tool function response")
911+
}
912+
return
913+
}
914+
}
915+
}
916+
917+
if !sawConfirmation {
918+
t.Fatal("expected confirmation event")
919+
}
920+
}
921+
868922
// Mock types for TArgs and TResults
869923
type TestArgs struct {
870924
Name string

‎tool/mcptoolset/set_test.go‎

Lines changed: 21 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -415,6 +415,9 @@ func TestMCPToolSetConfirmation(t *testing.T) {
415415
city: "Lisbon",
416416
want: []*genai.Content{
417417
genai.NewContentFromFunctionCall(toolName, map[string]any{"city": "Lisbon"}, "model"),
418+
genai.NewContentFromFunctionResponse(toolName, map[string]any{
419+
"error": errors.New("error tool \"get_weather\" requires confirmation, please approve or reject"),
420+
}, "user"),
418421
genai.NewContentFromFunctionCall(toolconfirmation.FunctionCallName, map[string]any{
419422
"originalFunctionCall": &genai.FunctionCall{
420423
Args: map[string]any{"city": "Lisbon"},
@@ -424,9 +427,6 @@ func TestMCPToolSetConfirmation(t *testing.T) {
424427
Hint: "Please approve or reject the tool call get_weather() by responding with a FunctionResponse with an expected ToolConfirmation payload.",
425428
},
426429
}, "model"),
427-
genai.NewContentFromFunctionResponse(toolName, map[string]any{
428-
"error": errors.New("error tool \"get_weather\" requires confirmation, please approve or reject"),
429-
}, "user"),
430430
},
431431
},
432432
{
@@ -438,6 +438,9 @@ func TestMCPToolSetConfirmation(t *testing.T) {
438438
confirmFunctionResponse: &genai.FunctionResponse{Name: toolconfirmation.FunctionCallName, Response: map[string]any{"confirmed": true}},
439439
want: []*genai.Content{
440440
genai.NewContentFromFunctionCall(toolName, map[string]any{"city": "Lisbon"}, "model"),
441+
genai.NewContentFromFunctionResponse(toolName, map[string]any{
442+
"error": errors.New("error tool \"get_weather\" requires confirmation, please approve or reject"),
443+
}, "user"),
441444
genai.NewContentFromFunctionCall(toolconfirmation.FunctionCallName, map[string]any{
442445
"originalFunctionCall": &genai.FunctionCall{
443446
Args: map[string]any{"city": "Lisbon"},
@@ -447,9 +450,6 @@ func TestMCPToolSetConfirmation(t *testing.T) {
447450
Hint: "Please approve or reject the tool call get_weather() by responding with a FunctionResponse with an expected ToolConfirmation payload.",
448451
},
449452
}, "model"),
450-
genai.NewContentFromFunctionResponse(toolName, map[string]any{
451-
"error": errors.New("error tool \"get_weather\" requires confirmation, please approve or reject"),
452-
}, "user"),
453453
genai.NewContentFromFunctionResponse(toolName, map[string]any{
454454
"output": map[string]any{"weather_summary": string(`Today in "Lisbon" is sunny`)},
455455
}, "user"),
@@ -465,6 +465,9 @@ func TestMCPToolSetConfirmation(t *testing.T) {
465465
confirmFunctionResponse: &genai.FunctionResponse{Name: toolconfirmation.FunctionCallName, Response: map[string]any{"confirmed": false}},
466466
want: []*genai.Content{
467467
genai.NewContentFromFunctionCall(toolName, map[string]any{"city": "Lisbon"}, "model"),
468+
genai.NewContentFromFunctionResponse(toolName, map[string]any{
469+
"error": errors.New("error tool \"get_weather\" requires confirmation, please approve or reject"),
470+
}, "user"),
468471
genai.NewContentFromFunctionCall(toolconfirmation.FunctionCallName, map[string]any{
469472
"originalFunctionCall": &genai.FunctionCall{
470473
Args: map[string]any{"city": "Lisbon"},
@@ -474,9 +477,6 @@ func TestMCPToolSetConfirmation(t *testing.T) {
474477
Hint: "Please approve or reject the tool call get_weather() by responding with a FunctionResponse with an expected ToolConfirmation payload.",
475478
},
476479
}, "model"),
477-
genai.NewContentFromFunctionResponse(toolName, map[string]any{
478-
"error": errors.New("error tool \"get_weather\" requires confirmation, please approve or reject"),
479-
}, "user"),
480480
genai.NewContentFromFunctionResponse(toolName, map[string]any{
481481
"error": errors.New("error tool \"get_weather\" call is rejected"),
482482
}, "user"),
@@ -523,6 +523,9 @@ func TestMCPToolSetConfirmation(t *testing.T) {
523523
city: "Lisbon",
524524
want: []*genai.Content{
525525
genai.NewContentFromFunctionCall(toolName, map[string]any{"city": "Lisbon"}, "model"),
526+
genai.NewContentFromFunctionResponse(toolName, map[string]any{
527+
"error": errors.New("error tool \"get_weather\" requires confirmation, please approve or reject"),
528+
}, "user"),
526529
genai.NewContentFromFunctionCall(toolconfirmation.FunctionCallName, map[string]any{
527530
"originalFunctionCall": &genai.FunctionCall{
528531
Args: map[string]any{"city": "Lisbon"},
@@ -532,9 +535,6 @@ func TestMCPToolSetConfirmation(t *testing.T) {
532535
Hint: "Please approve or reject the tool call get_weather() by responding with a FunctionResponse with an expected ToolConfirmation payload.",
533536
},
534537
}, "model"),
535-
genai.NewContentFromFunctionResponse(toolName, map[string]any{
536-
"error": errors.New("error tool \"get_weather\" requires confirmation, please approve or reject"),
537-
}, "user"),
538538
},
539539
},
540540
{
@@ -545,6 +545,9 @@ func TestMCPToolSetConfirmation(t *testing.T) {
545545
city: "Lisbon",
546546
want: []*genai.Content{
547547
genai.NewContentFromFunctionCall(toolName, map[string]any{"city": "Lisbon"}, "model"),
548+
genai.NewContentFromFunctionResponse(toolName, map[string]any{
549+
"error": errors.New("error tool \"get_weather\" requires confirmation, please approve or reject"),
550+
}, "user"),
548551
genai.NewContentFromFunctionCall(toolconfirmation.FunctionCallName, map[string]any{
549552
"originalFunctionCall": &genai.FunctionCall{
550553
Args: map[string]any{"city": "Lisbon"},
@@ -554,9 +557,6 @@ func TestMCPToolSetConfirmation(t *testing.T) {
554557
Hint: "Please approve or reject the tool call get_weather() by responding with a FunctionResponse with an expected ToolConfirmation payload.",
555558
},
556559
}, "model"),
557-
genai.NewContentFromFunctionResponse(toolName, map[string]any{
558-
"error": errors.New("error tool \"get_weather\" requires confirmation, please approve or reject"),
559-
}, "user"),
560560
},
561561
},
562562
{
@@ -568,6 +568,9 @@ func TestMCPToolSetConfirmation(t *testing.T) {
568568
confirmFunctionResponse: &genai.FunctionResponse{Name: toolconfirmation.FunctionCallName, Response: map[string]any{"confirmed": true}},
569569
want: []*genai.Content{
570570
genai.NewContentFromFunctionCall(toolName, map[string]any{"city": "Lisbon"}, "model"),
571+
genai.NewContentFromFunctionResponse(toolName, map[string]any{
572+
"error": errors.New("error tool \"get_weather\" requires confirmation, please approve or reject"),
573+
}, "user"),
571574
genai.NewContentFromFunctionCall(toolconfirmation.FunctionCallName, map[string]any{
572575
"originalFunctionCall": &genai.FunctionCall{
573576
Args: map[string]any{"city": "Lisbon"},
@@ -577,9 +580,6 @@ func TestMCPToolSetConfirmation(t *testing.T) {
577580
Hint: "Please approve or reject the tool call get_weather() by responding with a FunctionResponse with an expected ToolConfirmation payload.",
578581
},
579582
}, "model"),
580-
genai.NewContentFromFunctionResponse(toolName, map[string]any{
581-
"error": errors.New("error tool \"get_weather\" requires confirmation, please approve or reject"),
582-
}, "user"),
583583
genai.NewContentFromFunctionResponse(toolName, map[string]any{
584584
"output": map[string]any{"weather_summary": string(`Today in "Lisbon" is sunny`)},
585585
}, "user"),
@@ -595,6 +595,9 @@ func TestMCPToolSetConfirmation(t *testing.T) {
595595
confirmFunctionResponse: &genai.FunctionResponse{Name: toolconfirmation.FunctionCallName, Response: map[string]any{"confirmed": false}},
596596
want: []*genai.Content{
597597
genai.NewContentFromFunctionCall(toolName, map[string]any{"city": "Lisbon"}, "model"),
598+
genai.NewContentFromFunctionResponse(toolName, map[string]any{
599+
"error": errors.New("error tool \"get_weather\" requires confirmation, please approve or reject"),
600+
}, "user"),
598601
genai.NewContentFromFunctionCall(toolconfirmation.FunctionCallName, map[string]any{
599602
"originalFunctionCall": &genai.FunctionCall{
600603
Args: map[string]any{"city": "Lisbon"},
@@ -604,9 +607,6 @@ func TestMCPToolSetConfirmation(t *testing.T) {
604607
Hint: "Please approve or reject the tool call get_weather() by responding with a FunctionResponse with an expected ToolConfirmation payload.",
605608
},
606609
}, "model"),
607-
genai.NewContentFromFunctionResponse(toolName, map[string]any{
608-
"error": errors.New("error tool \"get_weather\" requires confirmation, please approve or reject"),
609-
}, "user"),
610610
genai.NewContentFromFunctionResponse(toolName, map[string]any{
611611
"error": errors.New("error tool \"get_weather\" call is rejected"),
612612
}, "user"),

0 commit comments

Comments
 (0)