@@ -48,6 +48,7 @@ import (
48
48
"cdr.dev/slog/sloggers/slogtest"
49
49
50
50
"github.com/coder/coder/v2/agent"
51
+ "github.com/coder/coder/v2/agent/agentcontainers"
51
52
"github.com/coder/coder/v2/agent/agentssh"
52
53
"github.com/coder/coder/v2/agent/agenttest"
53
54
"github.com/coder/coder/v2/agent/proto"
@@ -60,9 +61,16 @@ import (
60
61
"github.com/coder/coder/v2/tailnet"
61
62
"github.com/coder/coder/v2/tailnet/tailnettest"
62
63
"github.com/coder/coder/v2/testutil"
64
+ "github.com/coder/quartz"
63
65
)
64
66
65
67
func TestMain (m * testing.M ) {
68
+ if os .Getenv ("CODER_TEST_RUN_SUB_AGENT_MAIN" )== "1" {
69
+ // If we're running as a subagent, we don't want to run the main tests.
70
+ // Instead, we just run the subagent tests.
71
+ exit := runSubAgentMain ()
72
+ os .Exit (exit )
73
+ }
66
74
goleak .VerifyTestMain (m ,testutil .GoleakOptions ... )
67
75
}
68
76
@@ -1930,6 +1938,9 @@ func TestAgent_ReconnectingPTYContainer(t *testing.T) {
1930
1938
if os .Getenv ("CODER_TEST_USE_DOCKER" )!= "1" {
1931
1939
t .Skip ("Set CODER_TEST_USE_DOCKER=1 to run this test" )
1932
1940
}
1941
+ if _ ,err := exec .LookPath ("devcontainer" );err != nil {
1942
+ t .Skip ("This test requires the devcontainer CLI: npm install -g @devcontainers/cli" )
1943
+ }
1933
1944
1934
1945
pool ,err := dockertest .NewPool ("" )
1935
1946
require .NoError (t ,err ,"Could not connect to docker" )
@@ -1986,6 +1997,60 @@ func TestAgent_ReconnectingPTYContainer(t *testing.T) {
1986
1997
require .ErrorIs (t ,tr .ReadUntil (ctx ,nil ),io .EOF )
1987
1998
}
1988
1999
2000
+ type subAgentRequestPayload struct {
2001
+ Token string `json:"token"`
2002
+ Directory string `json:"directory"`
2003
+ }
2004
+
2005
+ // runSubAgentMain is the main function for the sub-agent that connects
2006
+ // to the control plane. It reads the CODER_AGENT_URL and
2007
+ // CODER_AGENT_TOKEN environment variables, sends the token, and exits
2008
+ // with a status code based on the response.
2009
+ func runSubAgentMain ()int {
2010
+ url := os .Getenv ("CODER_AGENT_URL" )
2011
+ token := os .Getenv ("CODER_AGENT_TOKEN" )
2012
+ if url == "" || token == "" {
2013
+ _ ,_ = fmt .Fprintln (os .Stderr ,"CODER_AGENT_URL and CODER_AGENT_TOKEN must be set" )
2014
+ return 10
2015
+ }
2016
+
2017
+ dir ,err := os .Getwd ()
2018
+ if err != nil {
2019
+ _ ,_ = fmt .Fprintf (os .Stderr ,"failed to get current working directory: %v\n " ,err )
2020
+ return 1
2021
+ }
2022
+ payload := subAgentRequestPayload {
2023
+ Token :token ,
2024
+ Directory :dir ,
2025
+ }
2026
+ b ,err := json .Marshal (payload )
2027
+ if err != nil {
2028
+ _ ,_ = fmt .Fprintf (os .Stderr ,"failed to marshal payload: %v\n " ,err )
2029
+ return 1
2030
+ }
2031
+
2032
+ req ,err := http .NewRequest ("POST" ,url ,bytes .NewReader (b ))
2033
+ if err != nil {
2034
+ _ ,_ = fmt .Fprintf (os .Stderr ,"failed to create request: %v\n " ,err )
2035
+ return 1
2036
+ }
2037
+ ctx ,cancel := context .WithTimeout (context .Background (),testutil .WaitLong )
2038
+ defer cancel ()
2039
+ req = req .WithContext (ctx )
2040
+ resp ,err := http .DefaultClient .Do (req )
2041
+ if err != nil {
2042
+ _ ,_ = fmt .Fprintf (os .Stderr ,"agent connection failed: %v\n " ,err )
2043
+ return 11
2044
+ }
2045
+ defer resp .Body .Close ()
2046
+ if resp .StatusCode != http .StatusOK {
2047
+ _ ,_ = fmt .Fprintf (os .Stderr ,"agent exiting with non-zero exit code %d\n " ,resp .StatusCode )
2048
+ return 12
2049
+ }
2050
+ _ ,_ = fmt .Println ("sub-agent connected successfully" )
2051
+ return 0
2052
+ }
2053
+
1989
2054
// This tests end-to-end functionality of auto-starting a devcontainer.
1990
2055
// It runs "devcontainer up" which creates a real Docker container. As
1991
2056
// such, it does not run by default in CI.
@@ -1999,6 +2064,56 @@ func TestAgent_DevcontainerAutostart(t *testing.T) {
1999
2064
if os .Getenv ("CODER_TEST_USE_DOCKER" )!= "1" {
2000
2065
t .Skip ("Set CODER_TEST_USE_DOCKER=1 to run this test" )
2001
2066
}
2067
+ if _ ,err := exec .LookPath ("devcontainer" );err != nil {
2068
+ t .Skip ("This test requires the devcontainer CLI: npm install -g @devcontainers/cli" )
2069
+ }
2070
+
2071
+ // This HTTP handler handles requests from runSubAgentMain which
2072
+ // acts as a fake sub-agent. We want to verify that the sub-agent
2073
+ // connects and sends its token. We use a channel to signal
2074
+ // that the sub-agent has connected successfully and then we wait
2075
+ // until we receive another signal to return from the handler. This
2076
+ // keeps the agent "alive" for as long as we want.
2077
+ subAgentConnected := make (chan subAgentRequestPayload ,1 )
2078
+ subAgentReady := make (chan struct {},1 )
2079
+ srv := httptest .NewServer (http .HandlerFunc (func (w http.ResponseWriter ,r * http.Request ) {
2080
+ t .Logf ("Sub-agent request received: %s %s" ,r .Method ,r .URL .Path )
2081
+
2082
+ if r .Method != http .MethodPost {
2083
+ http .Error (w ,"Method not allowed" ,http .StatusMethodNotAllowed )
2084
+ return
2085
+ }
2086
+
2087
+ // Read the token from the request body.
2088
+ var payload subAgentRequestPayload
2089
+ if err := json .NewDecoder (r .Body ).Decode (& payload );err != nil {
2090
+ http .Error (w ,"Failed to read token" ,http .StatusBadRequest )
2091
+ t .Logf ("Failed to read token: %v" ,err )
2092
+ return
2093
+ }
2094
+ defer r .Body .Close ()
2095
+
2096
+ t .Logf ("Sub-agent request payload received: %+v" ,payload )
2097
+
2098
+ // Signal that the sub-agent has connected successfully.
2099
+ select {
2100
+ case <- t .Context ().Done ():
2101
+ t .Logf ("Test context done, not processing sub-agent request" )
2102
+ return
2103
+ case subAgentConnected <- payload :
2104
+ }
2105
+
2106
+ // Wait for the signal to return from the handler.
2107
+ select {
2108
+ case <- t .Context ().Done ():
2109
+ t .Logf ("Test context done, not waiting for sub-agent ready" )
2110
+ return
2111
+ case <- subAgentReady :
2112
+ }
2113
+
2114
+ w .WriteHeader (http .StatusOK )
2115
+ }))
2116
+ defer srv .Close ()
2002
2117
2003
2118
pool ,err := dockertest .NewPool ("" )
2004
2119
require .NoError (t ,err ,"Could not connect to docker" )
@@ -2016,9 +2131,10 @@ func TestAgent_DevcontainerAutostart(t *testing.T) {
2016
2131
require .NoError (t ,err ,"create devcontainer directory" )
2017
2132
devcontainerFile := filepath .Join (devcontainerPath ,"devcontainer.json" )
2018
2133
err = os .WriteFile (devcontainerFile , []byte (`{
2019
- "name": "mywork",
2020
- "image": "busybox:latest",
2021
- "cmd": ["sleep", "infinity"]
2134
+ "name": "mywork",
2135
+ "image": "ubuntu:latest",
2136
+ "cmd": ["sleep", "infinity"],
2137
+ "runArgs": ["--network=host"]
2022
2138
}` ),0o600 )
2023
2139
require .NoError (t ,err ,"write devcontainer.json" )
2024
2140
@@ -2043,9 +2159,24 @@ func TestAgent_DevcontainerAutostart(t *testing.T) {
2043
2159
},
2044
2160
},
2045
2161
}
2162
+ mClock := quartz .NewMock (t )
2163
+ mClock .Set (time .Now ())
2164
+ tickerFuncTrap := mClock .Trap ().TickerFunc ("agentcontainers" )
2165
+
2046
2166
//nolint:dogsled
2047
- conn , _ ,_ ,_ ,_ := setupAgent (t ,manifest ,0 ,func (_ * agenttest.Client ,o * agent.Options ) {
2167
+ _ , agentClient ,_ ,_ ,_ := setupAgent (t ,manifest ,0 ,func (_ * agenttest.Client ,o * agent.Options ) {
2048
2168
o .ExperimentalDevcontainersEnabled = true
2169
+ o .ContainerAPIOptions = append (
2170
+ o .ContainerAPIOptions ,
2171
+ // Only match this specific dev container.
2172
+ agentcontainers .WithClock (mClock ),
2173
+ agentcontainers .WithContainerLabelIncludeFilter ("devcontainer.local_folder" ,tempWorkspaceFolder ),
2174
+ agentcontainers .WithSubAgentURL (srv .URL ),
2175
+ // The agent will copy "itself", but in the case of this test, the
2176
+ // agent is actually this test binary. So we'll tell the test binary
2177
+ // to execute the sub-agent main function via this env.
2178
+ agentcontainers .WithSubAgentEnv ("CODER_TEST_RUN_SUB_AGENT_MAIN=1" ),
2179
+ )
2049
2180
})
2050
2181
2051
2182
t .Logf ("Waiting for container with label: devcontainer.local_folder=%s" ,tempWorkspaceFolder )
@@ -2089,32 +2220,30 @@ func TestAgent_DevcontainerAutostart(t *testing.T) {
2089
2220
2090
2221
ctx := testutil .Context (t ,testutil .WaitLong )
2091
2222
2092
- ac ,err := conn .ReconnectingPTY (ctx ,uuid .New (),80 ,80 ,"" ,func (opts * workspacesdk.AgentReconnectingPTYInit ) {
2093
- opts .Container = container .ID
2094
- })
2095
- require .NoError (t ,err ,"failed to create ReconnectingPTY" )
2096
- defer ac .Close ()
2097
-
2098
- // Use terminal reader so we can see output in case somethin goes wrong.
2099
- tr := testutil .NewTerminalReader (t ,ac )
2223
+ // Ensure the container update routine runs.
2224
+ tickerFuncTrap .MustWait (ctx ).MustRelease (ctx )
2225
+ tickerFuncTrap .Close ()
2226
+ _ ,next := mClock .AdvanceNext ()
2227
+ next .MustWait (ctx )
2100
2228
2101
- require . NoError ( t , tr . ReadUntil ( ctx , func ( line string ) bool {
2102
- return strings . Contains ( line , "#" ) || strings . Contains ( line , "$" )
2103
- }), "find prompt " )
2229
+ // Verify that a subagent was created.
2230
+ subAgents := agentClient . GetSubAgents ( )
2231
+ require . Len ( t , subAgents , 1 , "expected one sub agent " )
2104
2232
2105
- wantFileName := "file-from-devcontainer"
2106
- wantFile := filepath .Join (tempWorkspaceFolder ,wantFileName )
2233
+ subAgent := subAgents [0 ]
2234
+ subAgentID ,err := uuid .FromBytes (subAgent .GetId ())
2235
+ require .NoError (t ,err ,"failed to parse sub-agent ID" )
2236
+ t .Logf ("Connecting to sub-agent: %s (ID: %s)" ,subAgent .Name ,subAgentID )
2107
2237
2108
- require .NoError (t ,json .NewEncoder (ac ).Encode (workspacesdk.ReconnectingPTYRequest {
2109
- // NOTE(mafredri): We must use absolute path here for some reason.
2110
- Data :fmt .Sprintf ("touch /workspaces/mywork/%s; exit\r " ,wantFileName ),
2111
- }),"create file inside devcontainer" )
2238
+ subAgentToken ,err := uuid .FromBytes (subAgent .GetAuthToken ())
2239
+ require .NoError (t ,err ,"failed to parse sub-agent token" )
2112
2240
2113
- // Wait for the connection to close to ensure the touch was executed.
2114
- require .ErrorIs (t ,tr .ReadUntil (ctx ,nil ),io .EOF )
2241
+ payload := testutil .RequireReceive (ctx ,t ,subAgentConnected )
2242
+ require .Equal (t ,subAgentToken .String (),payload .Token ,"sub-agent token should match" )
2243
+ require .Equal (t ,"/workspaces/mywork" ,payload .Directory ,"sub-agent directory should match" )
2115
2244
2116
- _ , err = os . Stat ( wantFile )
2117
- require . NoError ( t , err , "file should exist outside devcontainer" )
2245
+ // Allow the subagent to exit.
2246
+ close ( subAgentReady )
2118
2247
}
2119
2248
2120
2249
// TestAgent_DevcontainerRecreate tests that RecreateDevcontainer