From 2e188e346a0a846f5b8987594308db20413fe82f Mon Sep 17 00:00:00 2001
From: Virgil <virgil@lethean.io>
Date: Wed, 1 Apr 2026 06:53:33 +0000
Subject: [PATCH] feat(agentci): add context-aware ssh command helper

Thread dispatch SSH subprocesses through the caller context so cancellation applies to ticket transfer, remote cleanup, and existence checks.

Co-Authored-By: Virgil <virgil@lethean.io>
---
 agentci/security.go            | 13 ++++++++++++-
 agentci/security_test.go       | 14 ++++++++++++++
 jobrunner/handlers/dispatch.go |  6 +++---
 3 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/agentci/security.go b/agentci/security.go
index 16ae397..2a0f777 100644
--- a/agentci/security.go
+++ b/agentci/security.go
@@ -3,6 +3,7 @@
 package agentci
 
 import (
+	"context"
 	strings "dappco.re/go/core/scm/internal/ax/stringsx"
 	exec "golang.org/x/sys/execabs"
 	"path"
@@ -146,7 +147,17 @@ func EscapeShellArg(arg string) string {
 // SecureSSHCommand creates an SSH exec.Cmd with strict host key checking and batch mode.
 // Usage: SecureSSHCommand(...)
 func SecureSSHCommand(host string, remoteCmd string) *exec.Cmd {
-	return exec.Command("ssh",
+	return SecureSSHCommandContext(context.Background(), host, remoteCmd)
+}
+
+// SecureSSHCommandContext creates an SSH exec.Cmd with strict host key checking and batch mode.
+// Usage: SecureSSHCommandContext(...)
+func SecureSSHCommandContext(ctx context.Context, host string, remoteCmd string) *exec.Cmd {
+	if ctx == nil {
+		ctx = context.Background()
+	}
+
+	return exec.CommandContext(ctx, "ssh",
 		"-o", "StrictHostKeyChecking=yes",
 		"-o", "BatchMode=yes",
 		"-o", "ConnectTimeout=10",
diff --git a/agentci/security_test.go b/agentci/security_test.go
index 7030ede..cd808b2 100644
--- a/agentci/security_test.go
+++ b/agentci/security_test.go
@@ -3,6 +3,7 @@
 package agentci
 
 import (
+	"context"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -89,6 +90,19 @@ func TestSecureSSHCommand_Good(t *testing.T) {
 	assert.Equal(t, "ls -la", args[len(args)-1])
 }
 
+func TestSecureSSHCommandContext_Good(t *testing.T) {
+	cmd := SecureSSHCommandContext(context.Background(), "host.example.com", "ls -la")
+	args := cmd.Args
+
+	assert.Equal(t, "ssh", args[0])
+	assert.Contains(t, args, "-o")
+	assert.Contains(t, args, "StrictHostKeyChecking=yes")
+	assert.Contains(t, args, "BatchMode=yes")
+	assert.Contains(t, args, "ConnectTimeout=10")
+	assert.Equal(t, "host.example.com", args[len(args)-2])
+	assert.Equal(t, "ls -la", args[len(args)-1])
+}
+
 func TestMaskToken_Good(t *testing.T) {
 	tests := []struct {
 		name  string
diff --git a/jobrunner/handlers/dispatch.go b/jobrunner/handlers/dispatch.go
index f80be2b..166910c 100644
--- a/jobrunner/handlers/dispatch.go
+++ b/jobrunner/handlers/dispatch.go
@@ -296,7 +296,7 @@ func (h *DispatchHandler) secureTransfer(ctx context.Context, agent agentci.Agen
 	safePath := agentci.EscapeShellArg(remotePath)
 	remoteCmd := fmt.Sprintf("cat > %s && chmod %o %s", safePath, mode, safePath)
 
-	cmd := agentci.SecureSSHCommand(agent.Host, remoteCmd)
+	cmd := agentci.SecureSSHCommandContext(ctx, agent.Host, remoteCmd)
 	cmd.Stdin = bytes.NewReader(data)
 
 	output, err := cmd.CombinedOutput()
@@ -318,7 +318,7 @@ func (h *DispatchHandler) runRemote(ctx context.Context, agent agentci.AgentConf
 		remoteCmd = strings.Join(escaped, " ")
 	}
 
-	cmd := agentci.SecureSSHCommand(agent.Host, remoteCmd)
+	cmd := agentci.SecureSSHCommandContext(ctx, agent.Host, remoteCmd)
 	return cmd.Run()
 }
 
@@ -357,6 +357,6 @@ func (h *DispatchHandler) ticketExists(ctx context.Context, agent agentci.AgentC
 		"test -f %s || test -f %s || test -f %s",
 		queuePath, activePath, donePath,
 	)
-	cmd := agentci.SecureSSHCommand(agent.Host, checkCmd)
+	cmd := agentci.SecureSSHCommandContext(ctx, agent.Host, checkCmd)
 	return cmd.Run() == nil
 }