From e3fe96273adb61bf0a775f8d6b03b1e2e2d19f66 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Wed, 21 Feb 2024 01:51:22 -0800 Subject: [PATCH] nit Signed-off-by: Kevin Su --- flyteplugins/go/tasks/plugins/webapi/agent/plugin.go | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/flyteplugins/go/tasks/plugins/webapi/agent/plugin.go b/flyteplugins/go/tasks/plugins/webapi/agent/plugin.go index f30f5b203f..ea96980b78 100644 --- a/flyteplugins/go/tasks/plugins/webapi/agent/plugin.go +++ b/flyteplugins/go/tasks/plugins/webapi/agent/plugin.go @@ -5,7 +5,6 @@ import ( "encoding/gob" "fmt" "github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/service" - "io" "time" "golang.org/x/exp/maps" @@ -158,9 +157,12 @@ func (p Plugin) ExecuteTaskSync( } in, err := stream.Recv() - if err != nil && err != io.EOF { + if err != nil { return nil, nil, err } + if in.GetHeader() == nil { + return nil, nil, fmt.Errorf("expected header in the response, but got none") + } // TODO: Read the streaming output from the agent, and merge it into the final output. // For now, Propeller assumes that the output is always in the header. resource := in.GetHeader().GetResource()