feat: implement agent socket api, client and cli (#20758)

closes: https://github.com/coder/coder/issues/10352
closes: https://github.com/coder/internal/issues/1094
closes: https://github.com/coder/internal/issues/1095

In this pull request, we enable a new set of experimental cli commands
grouped under `coder exp sync`.
These commands allow any process acting within a coder workspace to
inform the coder agent of its requirements and execution progress. The
coder agent will then relay this information to other processes that
have subscribed.

These commands are:
```
# Check if this feature is enabled in your environment 
coder exp sync ping

# express that your unit depends on another
coder exp sync want <unit> <dependency_unit> 

# express that your unit intends to start a portion of the script that requires 
# other units to have completed first. This command blocks until all dependencies have been met
coder exp sync start <unit> 

# express that your unit has completes its work, allowing dependent units to begin their execution
coder exp sync complete <unit>
```

Example:

In order to automatically run claude code in a new workspace, it must
first have a git repository cloned. The scripts responsible for cloning
the repository and for running claude code would coordinate in the
following way:

```bash
# Script A: Claude code

# Inform the agent that the claude script wants the git script.
# That is, the git script must have completed before the claude script can begin its execution
coder exp sync want claude git

# Inform the agent that we would now like to begin execution of claude.
# This command will block until the git script (and any other defined dependencies)
# have completed
coder exp sync start claude

# Now we run claude code and any other commands we need
claude ...

# Once our script has completed, we inform the agent, so that any scripts that depend on this one
# may begin their execution

coder exp sync complete claude
```

```bash
# Script B: Git

# Because the git script does not have any dependencies, we can simply inform the agent that we 
# intend to start
coder exp sync start git

git clone ssh://git@github.com/coder/coder

# Once the repository have been cloned, we inform the agent that this script is complete, so that
# scripts that depend on it may begin their execution.
coder exp sync complete git
```

Notes:
* Unit names (ie. `claude` and `git`) given as input to the sync
commands are arbitrary strings. You do not have to conform to specific
identifiers. We recommend naming your scripts descriptively, but
succinctly.
* Scripts unit names should be well documented. Other scripts will need
to know the names you've chosen in order to depend on yours. Therefore,
you

---------

Co-authored-by: Mathias Fredriksson <mafredri@gmail.com>
This commit is contained in:
Sas Swart
2025-11-28 08:33:50 +02:00
committed by GitHub
parent ee58f40cad
commit ce627bf23f
37 changed files with 1314 additions and 277 deletions
+92 -173
View File
@@ -5,21 +5,18 @@ import (
"crypto/sha256"
"encoding/hex"
"fmt"
"net"
"os"
"path/filepath"
"runtime"
"testing"
"github.com/hashicorp/yamux"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"cdr.dev/slog"
"github.com/coder/coder/v2/agent/agentsocket"
"github.com/coder/coder/v2/agent/agentsocket/proto"
"github.com/coder/coder/v2/agent/unit"
"github.com/coder/coder/v2/codersdk/drpcsdk"
"github.com/coder/coder/v2/testutil"
)
// tempDirUnixSocket returns a temporary directory that can safely hold unix
@@ -47,23 +44,15 @@ func tempDirUnixSocket(t *testing.T) string {
}
// newSocketClient creates a DRPC client connected to the Unix socket at the given path.
func newSocketClient(t *testing.T, socketPath string) proto.DRPCAgentSocketClient {
func newSocketClient(ctx context.Context, t *testing.T, socketPath string) *agentsocket.Client {
t.Helper()
conn, err := net.Dial("unix", socketPath)
require.NoError(t, err)
config := yamux.DefaultConfig()
config.Logger = nil
session, err := yamux.Client(conn, config)
require.NoError(t, err)
client := proto.NewDRPCAgentSocketClient(drpcsdk.MultiplexedConn(session))
client, err := agentsocket.NewClient(ctx, agentsocket.WithPath(socketPath))
t.Cleanup(func() {
_ = session.Close()
_ = conn.Close()
_ = client.Close()
})
require.NoError(t, err)
return client
}
@@ -78,17 +67,17 @@ func TestDRPCAgentSocketService(t *testing.T) {
t.Parallel()
socketPath := filepath.Join(tempDirUnixSocket(t), "test.sock")
ctx := testutil.Context(t, testutil.WaitShort)
server, err := agentsocket.NewServer(
socketPath,
slog.Make().Leveled(slog.LevelDebug),
agentsocket.WithPath(socketPath),
)
require.NoError(t, err)
defer server.Close()
client := newSocketClient(t, socketPath)
client := newSocketClient(ctx, t, socketPath)
_, err = client.Ping(context.Background(), &proto.PingRequest{})
err = client.Ping(ctx)
require.NoError(t, err)
})
@@ -98,147 +87,116 @@ func TestDRPCAgentSocketService(t *testing.T) {
t.Run("NewUnit", func(t *testing.T) {
t.Parallel()
socketPath := filepath.Join(tempDirUnixSocket(t), "test.sock")
ctx := testutil.Context(t, testutil.WaitShort)
server, err := agentsocket.NewServer(
socketPath,
slog.Make().Leveled(slog.LevelDebug),
agentsocket.WithPath(socketPath),
)
require.NoError(t, err)
defer server.Close()
client := newSocketClient(t, socketPath)
client := newSocketClient(ctx, t, socketPath)
_, err = client.SyncStart(context.Background(), &proto.SyncStartRequest{
Unit: "test-unit",
})
err = client.SyncStart(ctx, "test-unit")
require.NoError(t, err)
status, err := client.SyncStatus(context.Background(), &proto.SyncStatusRequest{
Unit: "test-unit",
})
status, err := client.SyncStatus(ctx, "test-unit")
require.NoError(t, err)
require.Equal(t, "started", status.Status)
require.Equal(t, unit.StatusStarted, status.Status)
})
t.Run("UnitAlreadyStarted", func(t *testing.T) {
t.Parallel()
socketPath := filepath.Join(tempDirUnixSocket(t), "test.sock")
ctx := testutil.Context(t, testutil.WaitShort)
server, err := agentsocket.NewServer(
socketPath,
slog.Make().Leveled(slog.LevelDebug),
agentsocket.WithPath(socketPath),
)
require.NoError(t, err)
defer server.Close()
client := newSocketClient(t, socketPath)
client := newSocketClient(ctx, t, socketPath)
// First Start
_, err = client.SyncStart(context.Background(), &proto.SyncStartRequest{
Unit: "test-unit",
})
err = client.SyncStart(ctx, "test-unit")
require.NoError(t, err)
status, err := client.SyncStatus(context.Background(), &proto.SyncStatusRequest{
Unit: "test-unit",
})
status, err := client.SyncStatus(ctx, "test-unit")
require.NoError(t, err)
require.Equal(t, "started", status.Status)
require.Equal(t, unit.StatusStarted, status.Status)
// Second Start
_, err = client.SyncStart(context.Background(), &proto.SyncStartRequest{
Unit: "test-unit",
})
err = client.SyncStart(ctx, "test-unit")
require.ErrorContains(t, err, unit.ErrSameStatusAlreadySet.Error())
status, err = client.SyncStatus(context.Background(), &proto.SyncStatusRequest{
Unit: "test-unit",
})
status, err = client.SyncStatus(ctx, "test-unit")
require.NoError(t, err)
require.Equal(t, "started", status.Status)
require.Equal(t, unit.StatusStarted, status.Status)
})
t.Run("UnitAlreadyCompleted", func(t *testing.T) {
t.Parallel()
socketPath := filepath.Join(tempDirUnixSocket(t), "test.sock")
ctx := testutil.Context(t, testutil.WaitShort)
server, err := agentsocket.NewServer(
socketPath,
slog.Make().Leveled(slog.LevelDebug),
agentsocket.WithPath(socketPath),
)
require.NoError(t, err)
defer server.Close()
client := newSocketClient(t, socketPath)
client := newSocketClient(ctx, t, socketPath)
// First start
_, err = client.SyncStart(context.Background(), &proto.SyncStartRequest{
Unit: "test-unit",
})
err = client.SyncStart(ctx, "test-unit")
require.NoError(t, err)
status, err := client.SyncStatus(context.Background(), &proto.SyncStatusRequest{
Unit: "test-unit",
})
status, err := client.SyncStatus(ctx, "test-unit")
require.NoError(t, err)
require.Equal(t, "started", status.Status)
require.Equal(t, unit.StatusStarted, status.Status)
// Complete the unit
_, err = client.SyncComplete(context.Background(), &proto.SyncCompleteRequest{
Unit: "test-unit",
})
err = client.SyncComplete(ctx, "test-unit")
require.NoError(t, err)
status, err = client.SyncStatus(context.Background(), &proto.SyncStatusRequest{
Unit: "test-unit",
})
status, err = client.SyncStatus(ctx, "test-unit")
require.NoError(t, err)
require.Equal(t, "completed", status.Status)
require.Equal(t, unit.StatusComplete, status.Status)
// Second start
_, err = client.SyncStart(context.Background(), &proto.SyncStartRequest{
Unit: "test-unit",
})
err = client.SyncStart(ctx, "test-unit")
require.NoError(t, err)
status, err = client.SyncStatus(context.Background(), &proto.SyncStatusRequest{
Unit: "test-unit",
})
status, err = client.SyncStatus(ctx, "test-unit")
require.NoError(t, err)
require.Equal(t, "started", status.Status)
require.Equal(t, unit.StatusStarted, status.Status)
})
t.Run("UnitNotReady", func(t *testing.T) {
t.Parallel()
socketPath := filepath.Join(tempDirUnixSocket(t), "test.sock")
ctx := testutil.Context(t, testutil.WaitShort)
server, err := agentsocket.NewServer(
socketPath,
slog.Make().Leveled(slog.LevelDebug),
agentsocket.WithPath(socketPath),
)
require.NoError(t, err)
defer server.Close()
client := newSocketClient(t, socketPath)
client := newSocketClient(ctx, t, socketPath)
_, err = client.SyncWant(context.Background(), &proto.SyncWantRequest{
Unit: "test-unit",
DependsOn: "dependency-unit",
})
err = client.SyncWant(ctx, "test-unit", "dependency-unit")
require.NoError(t, err)
_, err = client.SyncStart(context.Background(), &proto.SyncStartRequest{
Unit: "test-unit",
})
err = client.SyncStart(ctx, "test-unit")
require.ErrorContains(t, err, "unit not ready")
status, err := client.SyncStatus(context.Background(), &proto.SyncStatusRequest{
Unit: "test-unit",
})
status, err := client.SyncStatus(ctx, "test-unit")
require.NoError(t, err)
require.Equal(t, string(unit.StatusPending), status.Status)
require.Equal(t, unit.StatusPending, status.Status)
require.False(t, status.IsReady)
})
})
@@ -250,107 +208,86 @@ func TestDRPCAgentSocketService(t *testing.T) {
t.Parallel()
socketPath := filepath.Join(tempDirUnixSocket(t), "test.sock")
ctx := testutil.Context(t, testutil.WaitShort)
server, err := agentsocket.NewServer(
socketPath,
slog.Make().Leveled(slog.LevelDebug),
agentsocket.WithPath(socketPath),
)
require.NoError(t, err)
defer server.Close()
client := newSocketClient(t, socketPath)
client := newSocketClient(ctx, t, socketPath)
// If dependency units are not registered, they are registered automatically
_, err = client.SyncWant(context.Background(), &proto.SyncWantRequest{
Unit: "test-unit",
DependsOn: "dependency-unit",
})
err = client.SyncWant(ctx, "test-unit", "dependency-unit")
require.NoError(t, err)
status, err := client.SyncStatus(context.Background(), &proto.SyncStatusRequest{
Unit: "test-unit",
})
status, err := client.SyncStatus(ctx, "test-unit")
require.NoError(t, err)
require.Len(t, status.Dependencies, 1)
require.Equal(t, "dependency-unit", status.Dependencies[0].DependsOn)
require.Equal(t, "completed", status.Dependencies[0].RequiredStatus)
require.Equal(t, unit.ID("dependency-unit"), status.Dependencies[0].DependsOn)
require.Equal(t, unit.StatusComplete, status.Dependencies[0].RequiredStatus)
})
t.Run("DependencyAlreadyRegistered", func(t *testing.T) {
t.Parallel()
socketPath := filepath.Join(tempDirUnixSocket(t), "test.sock")
ctx := testutil.Context(t, testutil.WaitShort)
server, err := agentsocket.NewServer(
socketPath,
slog.Make().Leveled(slog.LevelDebug),
agentsocket.WithPath(socketPath),
)
require.NoError(t, err)
defer server.Close()
client := newSocketClient(t, socketPath)
client := newSocketClient(ctx, t, socketPath)
// Start the dependency unit
_, err = client.SyncStart(context.Background(), &proto.SyncStartRequest{
Unit: "dependency-unit",
})
err = client.SyncStart(ctx, "dependency-unit")
require.NoError(t, err)
status, err := client.SyncStatus(context.Background(), &proto.SyncStatusRequest{
Unit: "dependency-unit",
})
status, err := client.SyncStatus(ctx, "dependency-unit")
require.NoError(t, err)
require.Equal(t, "started", status.Status)
require.Equal(t, unit.StatusStarted, status.Status)
// Add the dependency after the dependency unit has already started
_, err = client.SyncWant(context.Background(), &proto.SyncWantRequest{
Unit: "test-unit",
DependsOn: "dependency-unit",
})
err = client.SyncWant(ctx, "test-unit", "dependency-unit")
// Dependencies can be added even if the dependency unit has already started
require.NoError(t, err)
// The dependency is now reflected in the test unit's status
status, err = client.SyncStatus(context.Background(), &proto.SyncStatusRequest{
Unit: "test-unit",
})
status, err = client.SyncStatus(ctx, "test-unit")
require.NoError(t, err)
require.Equal(t, "dependency-unit", status.Dependencies[0].DependsOn)
require.Equal(t, "completed", status.Dependencies[0].RequiredStatus)
require.Equal(t, unit.ID("dependency-unit"), status.Dependencies[0].DependsOn)
require.Equal(t, unit.StatusComplete, status.Dependencies[0].RequiredStatus)
})
t.Run("DependencyAddedAfterDependentStarted", func(t *testing.T) {
t.Parallel()
socketPath := filepath.Join(tempDirUnixSocket(t), "test.sock")
ctx := testutil.Context(t, testutil.WaitShort)
server, err := agentsocket.NewServer(
socketPath,
slog.Make().Leveled(slog.LevelDebug),
agentsocket.WithPath(socketPath),
)
require.NoError(t, err)
defer server.Close()
client := newSocketClient(t, socketPath)
client := newSocketClient(ctx, t, socketPath)
// Start the dependent unit
_, err = client.SyncStart(context.Background(), &proto.SyncStartRequest{
Unit: "test-unit",
})
err = client.SyncStart(ctx, "test-unit")
require.NoError(t, err)
status, err := client.SyncStatus(context.Background(), &proto.SyncStatusRequest{
Unit: "test-unit",
})
status, err := client.SyncStatus(ctx, "test-unit")
require.NoError(t, err)
require.Equal(t, "started", status.Status)
require.Equal(t, unit.StatusStarted, status.Status)
// Add the dependency after the dependency unit has already started
_, err = client.SyncWant(context.Background(), &proto.SyncWantRequest{
Unit: "test-unit",
DependsOn: "dependency-unit",
})
err = client.SyncWant(ctx, "test-unit", "dependency-unit")
// Dependencies can be added even if the dependent unit has already started.
// The dependency applies the next time a unit is started. The current status is not updated.
@@ -359,12 +296,10 @@ func TestDRPCAgentSocketService(t *testing.T) {
require.NoError(t, err)
// The dependency is now reflected in the test unit's status
status, err = client.SyncStatus(context.Background(), &proto.SyncStatusRequest{
Unit: "test-unit",
})
status, err = client.SyncStatus(ctx, "test-unit")
require.NoError(t, err)
require.Equal(t, "dependency-unit", status.Dependencies[0].DependsOn)
require.Equal(t, "completed", status.Dependencies[0].RequiredStatus)
require.Equal(t, unit.ID("dependency-unit"), status.Dependencies[0].DependsOn)
require.Equal(t, unit.StatusComplete, status.Dependencies[0].RequiredStatus)
})
})
@@ -375,96 +310,80 @@ func TestDRPCAgentSocketService(t *testing.T) {
t.Parallel()
socketPath := filepath.Join(tempDirUnixSocket(t), "test.sock")
ctx := testutil.Context(t, testutil.WaitShort)
server, err := agentsocket.NewServer(
socketPath,
slog.Make().Leveled(slog.LevelDebug),
agentsocket.WithPath(socketPath),
)
require.NoError(t, err)
defer server.Close()
client := newSocketClient(t, socketPath)
client := newSocketClient(ctx, t, socketPath)
response, err := client.SyncReady(context.Background(), &proto.SyncReadyRequest{
Unit: "unregistered-unit",
})
ready, err := client.SyncReady(ctx, "unregistered-unit")
require.NoError(t, err)
require.False(t, response.Ready)
require.True(t, ready)
})
t.Run("UnitNotReady", func(t *testing.T) {
t.Parallel()
socketPath := filepath.Join(tempDirUnixSocket(t), "test.sock")
ctx := testutil.Context(t, testutil.WaitShort)
server, err := agentsocket.NewServer(
socketPath,
slog.Make().Leveled(slog.LevelDebug),
agentsocket.WithPath(socketPath),
)
require.NoError(t, err)
defer server.Close()
client := newSocketClient(t, socketPath)
client := newSocketClient(ctx, t, socketPath)
// Register a unit with an unsatisfied dependency
_, err = client.SyncWant(context.Background(), &proto.SyncWantRequest{
Unit: "test-unit",
DependsOn: "dependency-unit",
})
err = client.SyncWant(ctx, "test-unit", "dependency-unit")
require.NoError(t, err)
// Check readiness - should be false because dependency is not satisfied
response, err := client.SyncReady(context.Background(), &proto.SyncReadyRequest{
Unit: "test-unit",
})
ready, err := client.SyncReady(ctx, "test-unit")
require.NoError(t, err)
require.False(t, response.Ready)
require.False(t, ready)
})
t.Run("UnitReady", func(t *testing.T) {
t.Parallel()
socketPath := filepath.Join(tempDirUnixSocket(t), "test.sock")
ctx := testutil.Context(t, testutil.WaitShort)
server, err := agentsocket.NewServer(
socketPath,
slog.Make().Leveled(slog.LevelDebug),
agentsocket.WithPath(socketPath),
)
require.NoError(t, err)
defer server.Close()
client := newSocketClient(t, socketPath)
client := newSocketClient(ctx, t, socketPath)
// Register a unit with no dependencies - should be ready immediately
_, err = client.SyncStart(context.Background(), &proto.SyncStartRequest{
Unit: "test-unit",
})
err = client.SyncStart(ctx, "test-unit")
require.NoError(t, err)
// Check readiness - should be true
_, err = client.SyncReady(context.Background(), &proto.SyncReadyRequest{
Unit: "test-unit",
})
ready, err := client.SyncReady(ctx, "test-unit")
require.NoError(t, err)
require.True(t, ready)
// Also test a unit with satisfied dependencies
_, err = client.SyncWant(context.Background(), &proto.SyncWantRequest{
Unit: "dependent-unit",
DependsOn: "test-unit",
})
err = client.SyncWant(ctx, "dependent-unit", "test-unit")
require.NoError(t, err)
// Complete the dependency
_, err = client.SyncComplete(context.Background(), &proto.SyncCompleteRequest{
Unit: "test-unit",
})
err = client.SyncComplete(ctx, "test-unit")
require.NoError(t, err)
// Now dependent-unit should be ready
_, err = client.SyncReady(context.Background(), &proto.SyncReadyRequest{
Unit: "dependent-unit",
})
ready, err = client.SyncReady(ctx, "dependent-unit")
require.NoError(t, err)
require.True(t, ready)
})
})
}