Skip to content

Commit e9c8491

Browse files
committed
api/async exec: add schema refinement
1 parent 2797271 commit e9c8491

File tree

3 files changed

+141
-96
lines changed

3 files changed

+141
-96
lines changed

src/packages/backend/execute-code.test.ts

+28-3
Original file line numberDiff line numberDiff line change
@@ -92,13 +92,14 @@ describe("test timeout", () => {
9292

9393
describe("test env", () => {
9494
it("allows to specify environment variables", async () => {
95-
const { stdout, stderr } = await executeCode({
95+
const { stdout, stderr, type } = await executeCode({
9696
command: "sh",
9797
args: ["-c", "echo $FOO;"],
9898
err_on_exit: false,
9999
bash: false,
100100
env: { FOO: "bar" },
101101
});
102+
expect(type).toBe("blocking");
102103
expect(stdout).toBe("bar\n");
103104
expect(stderr).toBe("");
104105
});
@@ -142,17 +143,18 @@ describe("async", () => {
142143
expect(s.stdout).toEqual("foo\nbar\nbaz\n");
143144
expect(s.elapsed_s).toBeGreaterThan(0.1);
144145
expect(s.elapsed_s).toBeLessThan(3);
145-
expect(s.start).toBeGreaterThan(1);
146+
expect(s.start).toBeGreaterThan(Date.now() - 10 * 1000);
146147
expect(s.stderr).toEqual("");
147148
expect(s.exit_code).toEqual(0);
148149
}
149150
});
150151

151-
it("with an error", async () => {
152+
it("error/err_on_exit=true", async () => {
152153
const c = await executeCode({
153154
command: ">&2 echo baz; exit 3",
154155
bash: true,
155156
async_call: true,
157+
err_on_exit: true, // default
156158
});
157159
expect(c.type).toEqual("async");
158160
if (c.type !== "async") return;
@@ -170,6 +172,29 @@ describe("async", () => {
170172
expect(s.exit_code).toEqual(1);
171173
});
172174

175+
// without err_on_exit, the call is "completed" and we get the correct exit code
176+
it("error/err_on_exit=false", async () => {
177+
const c = await executeCode({
178+
command: ">&2 echo baz; exit 3",
179+
bash: true,
180+
async_call: true,
181+
err_on_exit: false,
182+
});
183+
expect(c.type).toEqual("async");
184+
if (c.type !== "async") return;
185+
const { job_id } = c;
186+
expect(typeof job_id).toEqual("string");
187+
if (typeof job_id !== "string") return;
188+
await new Promise((done) => setTimeout(done, 250));
189+
const s = await executeCode({ async_get: job_id });
190+
expect(s.type).toEqual("async");
191+
if (s.type !== "async") return;
192+
expect(s.status).toEqual("completed");
193+
expect(s.stdout).toEqual("");
194+
expect(s.stderr).toEqual("baz\n");
195+
expect(s.exit_code).toEqual(3);
196+
});
197+
173198
it("trigger a timeout", async () => {
174199
const c = await executeCode({
175200
command: "sh",

src/packages/backend/execute-code.ts

+1-1
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,7 @@ async function executeCodeNoAggregate(
8181
if (cached != null) {
8282
return cached;
8383
} else {
84-
throw new Error(`Async operation '${opts.async_get}' not found.`);
84+
throw new Error(`Async operation '${opts.async_get}' does not exist.`);
8585
}
8686
}
8787

src/packages/next/lib/api/schema/exec.ts

+112-92
Original file line numberDiff line numberDiff line change
@@ -4,113 +4,133 @@ import { FailedAPIOperationSchema } from "./common";
44
import { ComputeServerIdSchema } from "./compute/common";
55
import { ProjectIdSchema } from "./projects/common";
66

7-
// OpenAPI spec
8-
//
9-
export const ExecInputSchema = z
10-
.union([
11-
z.object({
12-
project_id: ProjectIdSchema,
13-
compute_server_id: ComputeServerIdSchema.describe(
14-
`If provided, the desired shell command will be run on the compute server whose id
7+
const ExecInputCommon = z.object({
8+
project_id: ProjectIdSchema,
9+
});
10+
11+
const ExecInputSchemaBlocking = ExecInputCommon.merge(
12+
z.object({
13+
compute_server_id: ComputeServerIdSchema.describe(
14+
`If provided, the desired shell command will be run on the compute server whose id
1515
is specified in this field (if available).`,
16-
).optional(),
17-
filesystem: z
18-
.boolean()
19-
.optional()
20-
.describe(
21-
`If \`true\`, this shell command runs in the fileserver container on the compute
16+
).optional(),
17+
filesystem: z
18+
.boolean()
19+
.optional()
20+
.describe(
21+
`If \`true\`, this shell command runs in the fileserver container on the compute
2222
server; otherwise, it runs on the main compute container.`,
23-
),
24-
path: z
25-
.string()
26-
.optional()
27-
.describe(
28-
"Path to working directory in which the shell command should be executed.",
29-
),
30-
command: z.string().describe("The shell command to execute."),
31-
args: z
32-
.array(z.string())
33-
.optional()
34-
.describe("An array of arguments to pass to the shell command."),
35-
timeout: z
36-
.number()
37-
.min(0)
38-
.default(60)
39-
.optional()
40-
.describe("Number of seconds before this shell command times out."),
41-
max_output: z
42-
.number()
43-
.min(0)
44-
.optional()
45-
.describe(
46-
"Maximum number of bytes to return from shell command output.",
47-
),
48-
bash: z
49-
.boolean()
50-
.optional()
51-
.describe(
52-
`If \`true\`, this command runs in a \`bash\` shell. To do so, the provided shell
23+
),
24+
path: z
25+
.string()
26+
.optional()
27+
.describe(
28+
"Path to working directory in which the shell command should be executed.",
29+
),
30+
command: z.string().describe("The shell command to execute."),
31+
args: z
32+
.array(z.string())
33+
.optional()
34+
.describe("An array of arguments to pass to the shell command."),
35+
timeout: z
36+
.number()
37+
.min(0)
38+
.default(60)
39+
.optional()
40+
.describe("Number of seconds before this shell command times out."),
41+
max_output: z
42+
.number()
43+
.min(0)
44+
.optional()
45+
.describe("Maximum number of bytes to return from shell command output."),
46+
bash: z
47+
.boolean()
48+
.optional()
49+
.describe(
50+
`If \`true\`, this command runs in a \`bash\` shell. To do so, the provided shell
5351
command is written to a file and then executed via the \`bash\` command.`,
54-
),
55-
home: z
56-
.string()
57-
.optional()
58-
.describe(
59-
`Specify \`$HOME\`. If not set, it is inferred from the environment's \`$HOME\``,
60-
),
61-
uid: z
62-
.number()
63-
.optional()
64-
.describe("Set the `UID` identity of the spawned process."),
65-
gid: z
66-
.number()
67-
.optional()
68-
.describe("Set the `GID` identity of the spawned process."),
69-
aggregate: z
70-
.union([
71-
z.number(),
72-
z.string(),
73-
z.object({ value: z.union([z.string(), z.number()]) }),
74-
])
75-
.optional()
76-
.describe(
77-
`If provided, this shell command is aggregated as in
52+
),
53+
home: z
54+
.string()
55+
.optional()
56+
.describe(
57+
`Specify \`$HOME\`. If not set, it is inferred from the environment's \`$HOME\``,
58+
),
59+
uid: z
60+
.number()
61+
.optional()
62+
.describe("Set the `UID` identity of the spawned process."),
63+
gid: z
64+
.number()
65+
.optional()
66+
.describe("Set the `GID` identity of the spawned process."),
67+
aggregate: z
68+
.union([
69+
z.number(),
70+
z.string(),
71+
z.object({ value: z.union([z.string(), z.number()]) }),
72+
])
73+
.optional()
74+
.describe(
75+
`If provided, this shell command is aggregated as in
7876
\`src/packages/backend/aggregate.js\`. This parameter allows one to specify
7977
multiple callbacks to be executed against the output of the same command
8078
(given identical arguments) within a 60-second window.`,
81-
),
82-
err_on_exit: z
83-
.boolean()
84-
.optional()
85-
.describe(
86-
`When \`true\` (the default),
79+
),
80+
err_on_exit: z
81+
.boolean()
82+
.optional()
83+
.describe(
84+
`When \`true\` (the default),
8785
this call will throw an error whenever the provided shell command
8886
exits with a non-zero exit code.`,
89-
),
90-
env: z
91-
.record(z.string(), z.string())
92-
.optional()
93-
.describe(
94-
"Environment variables to be passed to the shell command upon execution.",
95-
),
96-
async_call: z.boolean().optional()
97-
.describe(`If \`true\`, the execution happens asynchronously.
87+
),
88+
env: z
89+
.record(z.string(), z.string())
90+
.optional()
91+
.describe(
92+
"Environment variables to be passed to the shell command upon execution.",
93+
),
94+
async_call: z.boolean().optional()
95+
.describe(`If \`true\`, the execution happens asynchronously.
9896
The API call does not block and returns an ID (\`job_id\`).
97+
9998
Later, use that ID in a call to \`async_get\` to get status updates, partial output, and eventually the final result.
99+
In such a call, you also have to set the \`project_id\`, because the results are cached in the project.
100+
101+
Additionally and if not specified, \`max_output\` is set to 1MB and and \`timeout\` to 10 minutes.
100102
101-
This does not support executing code on compute servers – only inside the project itself.
103+
NOTE: This does not support executing code on compute servers – only inside the project itself.
102104
103-
Additionally and if not specified, \`max_output\` is set to 1MB and and \`timeout\` to 10 minutes.`),
104-
}),
105+
HINT: set \`err_on_exit=false\`, to recieve the real \`exit_code\` of the executed command and status ends with "completed", unless there is a fundamental problem running the command.
106+
`),
107+
}),
108+
);
105109

106-
z.object({
107-
project_id: ProjectIdSchema,
108-
async_get: z.string().optional()
109-
.describe(`For a given \`job_id\`, returned when setting \`async_call=true\`,
110+
const ExecInputSchemaAsync = ExecInputCommon.merge(
111+
z.object({
112+
project_id: ProjectIdSchema,
113+
async_get: z.string().optional()
114+
.describe(`For a given \`job_id\`, which has been returned when setting \`async_call=true\`,
110115
retrieve the corresponding status or the result.
116+
117+
The returned object contains the current \`stdout\` and \`stderr\` output,
118+
as well as a status field indicating if the job is still running or has completed.
119+
Start time and duration are returned as well.
120+
111121
Results are cached temporarily in the project.`),
112-
}),
113-
])
122+
}),
123+
);
124+
125+
export const ExecInputSchema = z
126+
.union([ExecInputSchemaBlocking, ExecInputSchemaAsync])
127+
.refine((data) => {
128+
if ("async_get" in data) {
129+
return ExecInputSchemaAsync.safeParse(data).success;
130+
} else {
131+
return ExecInputSchemaBlocking.safeParse(data).success;
132+
}
133+
})
114134
.describe("Perform arbitrary shell commands in a compute server or project.");
115135

116136
const ExecOutputBlocking = z.object({

0 commit comments

Comments
 (0)