From f4fb48f423469450751ef99c38fcaab0e225d7bc Mon Sep 17 00:00:00 2001
From: Jon Langevin <codingsloth@pm.me>
Date: Thu, 21 May 2026 01:52:38 -0400
Subject: [PATCH] fix(wfctl): add log follow completion grace

---
 cmd/wfctl/logs.go      |  4 ++-
 cmd/wfctl/logs_test.go | 60 ++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 61 insertions(+), 3 deletions(-)

diff --git a/cmd/wfctl/logs.go b/cmd/wfctl/logs.go
index 164afdc5..07854f86 100644
--- a/cmd/wfctl/logs.go
+++ b/cmd/wfctl/logs.go
@@ -13,6 +13,8 @@ import (
 	"github.com/GoCodeAlone/workflow/interfaces"
 )
 
+const logCaptureFollowCompletionGrace = 30 * time.Second
+
 func runLogs(args []string) error {
 	return runLogsWithOutput(args, os.Stdout)
 }
@@ -113,7 +115,7 @@ func runLogsCapture(args []string, out io.Writer) error {
 	ctx := context.Background()
 	if follow && duration > 0 {
 		var cancel context.CancelFunc
-		ctx, cancel = context.WithTimeout(ctx, duration)
+		ctx, cancel = context.WithTimeout(ctx, duration+logCaptureFollowCompletionGrace)
 		defer cancel()
 	}
 	durationSeconds := int64(0)
diff --git a/cmd/wfctl/logs_test.go b/cmd/wfctl/logs_test.go
index d824a3f5..8a2f9afa 100644
--- a/cmd/wfctl/logs_test.go
+++ b/cmd/wfctl/logs_test.go
@@ -8,6 +8,7 @@ import (
 	"path/filepath"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/GoCodeAlone/workflow/config"
 	"github.com/GoCodeAlone/workflow/interfaces"
@@ -15,11 +16,18 @@ import (
 
 type fakeLogProvider struct {
 	applyCapture
-	req interfaces.LogCaptureRequest
+	req          interfaces.LogCaptureRequest
+	ctxDeadline  time.Time
+	ctxRemaining time.Duration
+	hasDeadline  bool
 }
 
-func (p *fakeLogProvider) CaptureLogs(_ context.Context, req interfaces.LogCaptureRequest, sink interfaces.LogCaptureSink) error {
+func (p *fakeLogProvider) CaptureLogs(ctx context.Context, req interfaces.LogCaptureRequest, sink interfaces.LogCaptureSink) error {
 	p.req = req
+	p.ctxDeadline, p.hasDeadline = ctx.Deadline()
+	if p.hasDeadline {
+		p.ctxRemaining = time.Until(p.ctxDeadline)
+	}
 	return sink.WriteLogChunk(interfaces.LogChunk{Data: []byte("line one\n"), Source: "historic"})
 }
 
@@ -177,6 +185,54 @@ modules:
 	}
 }
 
+func TestLogsCaptureFollowContextAllowsProviderCompletionGrace(t *testing.T) {
+	tmp := t.TempDir()
+	cfg := filepath.Join(tmp, "app.yaml")
+	if err := os.WriteFile(cfg, []byte(`
+version: "1"
+modules:
+  - name: do
+    type: iac.provider
+    config:
+      provider: digitalocean
+  - name: web
+    type: infra.container_service
+    config:
+      provider: do
+      app_name: bmw-staging
+`), 0o600); err != nil {
+		t.Fatal(err)
+	}
+
+	provider := &fakeLogProvider{}
+	orig := resolveIaCProvider
+	resolveIaCProvider = func(_ context.Context, _ string, _ map[string]any) (interfaces.IaCProvider, io.Closer, error) {
+		return provider, nil, nil
+	}
+	t.Cleanup(func() { resolveIaCProvider = orig })
+
+	var out bytes.Buffer
+	err := runLogsWithOutput([]string{
+		"capture",
+		"--config", cfg,
+		"--resource", "web",
+		"--follow",
+		"--duration", "5s",
+	}, &out)
+	if err != nil {
+		t.Fatalf("runLogsWithOutput: %v", err)
+	}
+	if provider.req.DurationSeconds != 5 {
+		t.Fatalf("DurationSeconds = %d, want 5", provider.req.DurationSeconds)
+	}
+	if !provider.hasDeadline {
+		t.Fatal("expected wfctl to retain a client-side timeout guard")
+	}
+	if got := provider.ctxRemaining; got <= 5*time.Second {
+		t.Fatalf("client context remaining = %v, want grace beyond requested 5s follow", got)
+	}
+}
+
 func TestLogsCaptureRejectsUnknownType(t *testing.T) {
 	tmp := t.TempDir()
 	cfg := filepath.Join(tmp, "app.yaml")