internal/pluginhost/issue6332_delivery_test.go
Rama main · réplica de lectura del repositorio original.
package pluginhost
import (
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/MoeForge/moego/ojson"
"github.com/router-for-me/CLIProxyAPI/v8/internal/cache"
"github.com/router-for-me/CLIProxyAPI/v8/internal/config"
"github.com/router-for-me/CLIProxyAPI/v8/internal/registry"
runtimeexecutor "github.com/router-for-me/CLIProxyAPI/v8/internal/runtime/executor"
_ "github.com/router-for-me/CLIProxyAPI/v8/internal/translator"
"github.com/router-for-me/CLIProxyAPI/v8/sdk/api/handlers"
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth"
"github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/usage"
sdkconfig "github.com/router-for-me/CLIProxyAPI/v8/sdk/config"
"github.com/router-for-me/CLIProxyAPI/v8/sdk/pluginabi"
"github.com/router-for-me/CLIProxyAPI/v8/sdk/pluginapi"
)
type issue6332NestedLifecycle struct {
handlers.PluginInterceptorHost
completed chan pluginapi.RequestCompletion
}
func (*issue6332NestedLifecycle) HasRequestInterceptors() bool { return false }
func (*issue6332NestedLifecycle) HasStreamInterceptors() bool { return false }
func (h *issue6332NestedLifecycle) CompleteRequest(_ context.Context, c pluginapi.RequestCompletion) {
h.completed <- c
}
// Exercise the actual plugin RPC callbacks and stream bridge, not a fake model executor.
func TestIssue6332PluginNestedDelivery(t *testing.T) {
for _, mode := range []string{"complete", "cancel_partial", "cancel_complete"} {
t.Run(mode, func(t *testing.T) {
cancelStream := mode != "complete"
const model = "gemini-3.7-flash"
signature := strings.Repeat("s", 128)
requests := make(chan []byte, 2)
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
body, errRead := io.ReadAll(r.Body)
if errRead != nil {
t.Error(errRead)
return
}
requests <- body
w.Header().Set("Content-Type", "text/event-stream")
// Cancellation of partial upstream output must not commit replay.
finish := `,"finishReason":"STOP"`
if mode == "cancel_partial" {
finish = ""
}
_, _ = io.WriteString(w, `data: {"response":{"candidates":[{"content":{"role":"model","parts":[{"functionCall":{"id":"native-call","name":"lookup","args":{"key":"value"}},"thoughtSignature":"`+signature+`"}]}`+finish+`}],"usageMetadata":{"promptTokenCount":11,"candidatesTokenCount":22,"totalTokenCount":33},"responseId":"nested"}}`+"\n\n")
w.(http.Flusher).Flush()
if cancelStream {
<-r.Context().Done()
}
}))
defer server.Close()
manager := cliproxyauth.NewManager(nil, nil, nil)
manager.RegisterExecutor(runtimeexecutor.NewAntigravityExecutor(&config.Config{}))
auth := &cliproxyauth.Auth{ID: t.Name(), Provider: "antigravity", Status: cliproxyauth.StatusActive, ProxyURL: "direct", Metadata: map[string]any{"access_token": "token", "expired": time.Now().Add(time.Hour).Format(time.RFC3339), "project_id": "project"}, Attributes: map[string]string{"base_url": server.URL}}
if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil {
t.Fatal(errRegister)
}
registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: model}})
t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) })
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
lifecycle := &issue6332NestedLifecycle{completed: make(chan pluginapi.RequestCompletion, 4)}
base.SetPluginHost(lifecycle)
host := New()
host.SetModelExecutor(base)
outer, finishOuter := usage.WithStreamDelivery(context.Background())
// No outer completion is acknowledged until the nested tool loop finishes.
defer finishOuter(context.Canceled)
callbackID, closeCallback := host.openCallbackContextForPlugin(outer, "tool-loop-plugin")
defer closeCallback()
call := func(body string) (string, error) {
raw, _ := json.Marshal(rpcHostModelExecutionRequest{HostCallbackID: callbackID, HostModelExecutionRequest: pluginapi.HostModelExecutionRequest{EntryProtocol: "openai-response", ExitProtocol: "openai-response", Model: model, Stream: true, Body: []byte(body)}})
response, errCall := host.callFromPlugin(context.Background(), pluginabi.MethodHostModelExecuteStream, raw)
if errCall != nil {
return "", errCall
}
stream, errDecode := decodeRPCEnvelope[pluginapi.HostModelStreamResponse](response)
if errDecode != nil {
return "", errDecode
}
read, _ := json.Marshal(pluginapi.HostModelStreamReadRequest{StreamID: stream.StreamID})
var output strings.Builder
for {
response, errRead := host.callFromPlugin(context.Background(), pluginabi.MethodHostModelStreamRead, read)
if errRead != nil {
return "", errRead
}
chunk, errDecodeChunk := decodeRPCEnvelope[pluginapi.HostModelStreamReadResponse](response)
if errDecodeChunk != nil {
return "", errDecodeChunk
}
if chunk.Error != "" {
return "", fmt.Errorf("stream error: %s", chunk.Error)
}
output.Write(chunk.Payload)
if cancelStream {
closeCallback()
return output.String(), nil
}
if chunk.Done {
return output.String(), nil
}
}
}
result := make(chan error, 1)
go func() {
first := `{"model":"` + model + `","stream":true,"metadata":{"session_id":"` + t.Name() + `"},"input":[{"role":"user","content":"lookup"}]}`
output, errCall := call(first)
if errCall != nil || cancelStream {
result <- errCall
return
}
var callID string
for _, line := range strings.Split(output, "\n") {
event := ojson.Parse(strings.TrimSpace(strings.TrimPrefix(line, "data:")))
for _, item := range event.Get("response.output").Array() {
if item.Get("type").String() == "function_call" {
callID = item.Get("call_id").String()
}
}
}
if callID == "" {
result <- fmt.Errorf("missing function call: %s", output)
return
}
doc := ojson.NewDocString(first)
_ = doc.Set("input.-1", map[string]any{"type": "function_call", "call_id": callID, "name": "lookup", "arguments": `{"key":"value"}`})
_ = doc.Set("input.-1", map[string]any{"type": "function_call_output", "call_id": callID, "output": "found"})
second := doc.String()
_, errCall = call(second)
result <- errCall
}()
select {
case errResult := <-result:
if errResult != nil {
t.Fatal(errResult)
}
case <-time.After(3 * time.Second):
finishOuter(context.Canceled)
closeCallback()
<-result
t.Fatal("nested tool loop blocked waiting for outer HTTP delivery")
}
if usage.StreamDeliverySupported(outer) {
t.Error("nested producer leaked outer delivery support")
}
turns := 2
if cancelStream {
turns = 1
}
for range turns {
select {
case c := <-lifecycle.completed:
want := pluginapi.RequestCompletionSucceeded
if cancelStream {
want = pluginapi.RequestCompletionCanceled
}
if c.Outcome != want {
t.Errorf("nested lifecycle = %+v, want %s", c, want)
}
case <-time.After(3 * time.Second):
t.Fatal("nested lifecycle waits for outer delivery")
}
}
<-requests
if !cancelStream {
second := <-requests
if !strings.Contains(string(second), signature) {
t.Errorf("second turn lost replay signature: %s", second)
}
} else {
readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
defer cancelRead()
items, _, _, errGet := cache.GetAntigravityReasoningReplayItemsWithSnapshotRequired(readCtx, model, "responses:"+t.Name())
if errGet != nil || (len(items) > 0) != (mode == "cancel_complete") {
t.Errorf("replay must follow upstream completeness, not delivery: %v, %v", items, errGet)
}
}
})
}
}