From 279f8cd482f185727bb864710397d2edd85e40c0 Mon Sep 17 00:00:00 2001 From: Petr Date: Fri, 15 May 2026 08:30:47 +0200 Subject: [PATCH 01/12] fix(ui): drop hardcoded 'Petr' from Dashboard greeting (#285) Greeting now shows only the time-of-day phrase (e.g. 'Good Morning') instead of '${greeting()}, Petr'. The hardcoded name violated the no-hardcoded-defaults rule and was visible to every user. --- web/frontend/src/pages/Dashboard.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/web/frontend/src/pages/Dashboard.tsx b/web/frontend/src/pages/Dashboard.tsx index ad659cb3..f1adb39c 100644 --- a/web/frontend/src/pages/Dashboard.tsx +++ b/web/frontend/src/pages/Dashboard.tsx @@ -93,7 +93,7 @@ export function DashboardPage() { return (
Date: Fri, 15 May 2026 08:46:06 +0200 Subject: [PATCH 02/12] feat(projects): show organization name per project (#290) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Persist org_id + org_name on ProjectConfig and surface them in the Projects UI so users with multiple Keboola organizations can tell which org each project belongs to. Sources of org info, by precedence: 1. Storage API /v2/storage/tokens/verify owner.organization (used by both `project add` and `org setup` -- free when present). 2. Manage API list_organization_projects per-project organization payload (used during `org setup`). 3. Manage API GET /manage/organizations/{id} fallback (one call per `org setup` run when neither source above carries the name). ProjectConfig fields default to None; pre-existing config.json files load cleanly. Projects without org info render '—' in the new ORG column. --- src/keboola_agent_cli/client.py | 3 + src/keboola_agent_cli/manage_client.py | 18 ++++++ src/keboola_agent_cli/models.py | 16 +++++ src/keboola_agent_cli/services/org_service.py | 25 +++++++- .../services/project_service.py | 6 ++ tests/helpers.py | 4 ++ tests/test_client.py | 29 +++++++++ tests/test_manage_client.py | 19 ++++++ tests/test_models.py | 28 ++++++++ tests/test_org_service.py | 64 +++++++++++++++++++ tests/test_services.py | 38 +++++++++++ web/frontend/src/pages/Projects.tsx | 11 ++++ web/frontend/src/types.ts | 2 + 13 files changed, 262 insertions(+), 1 deletion(-) diff --git a/src/keboola_agent_cli/client.py b/src/keboola_agent_cli/client.py index e0a36ef2..9c545950 100644 --- a/src/keboola_agent_cli/client.py +++ b/src/keboola_agent_cli/client.py @@ -221,6 +221,7 @@ def verify_token(self) -> TokenVerifyResponse: data = response.json() owner = data.get("owner", {}) + org = owner.get("organization") or {} response = TokenVerifyResponse( token_id=str(data.get("id", "")), token_description=data.get("description", ""), @@ -229,6 +230,8 @@ def verify_token(self) -> TokenVerifyResponse: owner_name=owner.get("name", ""), default_backend=owner.get("defaultBackend", "snowflake"), features=owner.get("features", []), + org_id=org.get("id"), + org_name=org.get("name") or None, ) # Refresh the features cache on every successful verify so explicit # callers stay consistent with the cached view used by has_feature(). diff --git a/src/keboola_agent_cli/manage_client.py b/src/keboola_agent_cli/manage_client.py index fb391a97..4c16a9fe 100644 --- a/src/keboola_agent_cli/manage_client.py +++ b/src/keboola_agent_cli/manage_client.py @@ -76,6 +76,24 @@ def get_project(self, project_id: int) -> dict[str, Any]: response = self._do_request("GET", f"/manage/projects/{project_id}") return response.json() + def get_organization(self, org_id: int) -> dict[str, Any]: + """Get organization details by ID. + + Used to resolve the organization name from its ID (e.g. during + `org setup`, where only the org_id is known up front). + + Args: + org_id: The organization ID. + + Returns: + Organization dict with at least 'id' and 'name' fields. + + Raises: + KeboolaApiError: On API errors (e.g. 403 if not an org member). + """ + response = self._do_request("GET", f"/manage/organizations/{org_id}") + return response.json() + def list_organization_projects(self, org_id: int) -> list[dict[str, Any]]: """List all projects in an organization. diff --git a/src/keboola_agent_cli/models.py b/src/keboola_agent_cli/models.py index 024ac031..013cc767 100644 --- a/src/keboola_agent_cli/models.py +++ b/src/keboola_agent_cli/models.py @@ -20,6 +20,14 @@ class ProjectConfig(BaseModel): default=None, description="Active development branch ID (None = main/production branch)", ) + org_id: int | None = Field( + default=None, + description="Organization ID (populated via `org setup` or when verify_token returns it)", + ) + org_name: str | None = Field( + default=None, + description="Organization name (populated via `org setup` or when verify_token returns it)", + ) @field_validator("stack_url") @classmethod @@ -105,6 +113,14 @@ class TokenVerifyResponse(BaseModel): default_factory=list, description="Project feature flags (e.g. agent-chat, storage-types)", ) + org_id: int | None = Field( + default=None, + description="Organization ID parsed from owner.organization (when present)", + ) + org_name: str | None = Field( + default=None, + description="Organization name parsed from owner.organization (when present)", + ) class ComponentDetail(BaseModel): diff --git a/src/keboola_agent_cli/services/org_service.py b/src/keboola_agent_cli/services/org_service.py index 33c3e686..0746e206 100644 --- a/src/keboola_agent_cli/services/org_service.py +++ b/src/keboola_agent_cli/services/org_service.py @@ -108,6 +108,7 @@ def setup_organization( raise ValueError(msg) manage_client = self._manage_client_factory(stack_url, manage_token) + org_name: str | None = None try: if project_ids: projects, fetch_failed = self._fetch_projects_by_ids(manage_client, project_ids) @@ -118,6 +119,20 @@ def setup_organization( projects = manage_client.list_organization_projects(org_id) fetch_failed = [] + # Derive org_name: prefer per-project payload, otherwise call Manage API once. + for p in projects: + candidate = (p.get("organization") or {}).get("name") + if isinstance(candidate, str) and candidate: + org_name = candidate + break + if not org_name and org_id: + try: + fetched = manage_client.get_organization(org_id).get("name") + if isinstance(fetched, str) and fetched: + org_name = fetched + except Exception: + logger.debug("Could not resolve organization name for org_id=%s", org_id) + # Resolve token owner identity for unique token naming owner_name = "" try: @@ -183,6 +198,8 @@ def setup_organization( token_description=token_description, owner_name=owner_name, token_expires_in=token_expires_in, + org_id=org_id, + org_name=org_name, ) # Re-read to get masked token registered = self._config_store.get_project(alias) @@ -517,6 +534,8 @@ def _setup_single_project( token_description: str, owner_name: str = "", token_expires_in: int | None = None, + org_id: int | None = None, + org_name: str | None = None, ) -> None: """Create a token for a single project, verify it, and register it. @@ -559,12 +578,16 @@ def _setup_single_project( finally: storage_client.close() - # Register the project in config + # Register the project in config. Org info from token verify wins + # (it's the authoritative source); fall back to the values derived + # during org setup. project_config = ProjectConfig( stack_url=stack_url, token=storage_token, project_name=token_info.project_name, project_id=token_info.project_id, + org_id=token_info.org_id or org_id, + org_name=token_info.org_name or org_name, ) self._config_store.add_project(alias, project_config) diff --git a/src/keboola_agent_cli/services/project_service.py b/src/keboola_agent_cli/services/project_service.py index 468f95bb..30fc4348 100644 --- a/src/keboola_agent_cli/services/project_service.py +++ b/src/keboola_agent_cli/services/project_service.py @@ -62,6 +62,8 @@ def add_project(self, alias: str, stack_url: str, token: str) -> dict[str, Any]: token=token, project_name=token_info.project_name, project_id=token_info.project_id, + org_id=token_info.org_id, + org_name=token_info.org_name, ) self._config_store.add_project(alias, project) @@ -72,6 +74,8 @@ def add_project(self, alias: str, stack_url: str, token: str) -> dict[str, Any]: "project_id": token_info.project_id, "stack_url": stack_url, "token": mask_token(token), + "org_id": token_info.org_id, + "org_name": token_info.org_name, } def remove_project(self, alias: str) -> dict[str, str]: @@ -588,6 +592,8 @@ def list_projects(self) -> list[dict[str, Any]]: "token": mask_token(project.token), "is_default": alias == config.default_project, "active_branch_id": project.active_branch_id, + "org_id": project.org_id, + "org_name": project.org_name, } ) return result diff --git a/tests/helpers.py b/tests/helpers.py index a7fbc5cf..440c6c85 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -16,6 +16,8 @@ def make_mock_client( project_name: str = "Test Project", project_id: int = 1234, token_description: str = "My Token", + org_id: int | None = None, + org_name: str | None = None, ) -> MagicMock: """Create a mock KeboolaClient that returns a successful verify_token response. @@ -29,6 +31,8 @@ def make_mock_client( project_id=project_id, project_name=project_name, owner_name=project_name, + org_id=org_id, + org_name=org_name, ) return mock_client diff --git a/tests/test_client.py b/tests/test_client.py index a8a163b5..3e4107ad 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -53,6 +53,35 @@ def test_verify_token_success(self, httpx_mock) -> None: assert result.project_id == 1234 assert result.token_description == "My test token" assert result.token_id == "12345" + # Owner without organization sub-object yields None org fields. + assert result.org_id is None + assert result.org_name is None + client.close() + + def test_verify_token_extracts_organization(self, httpx_mock) -> None: + """When owner.organization is present, org_id and org_name are parsed.""" + httpx_mock.add_response( + url="https://connection.keboola.com/v2/storage/tokens/verify", + json={ + "id": "12345", + "description": "tok", + "owner": { + "id": 1234, + "name": "Test Project", + "organization": {"id": 438, "name": "Keboola Demo"}, + }, + }, + status_code=200, + ) + + client = KeboolaClient( + stack_url="https://connection.keboola.com", + token="901-10493007-VDtlEDWDF6Tx5V8jjE8FshFlqM0Hl0c08KHqpt0k", + ) + result = client.verify_token() + + assert result.org_id == 438 + assert result.org_name == "Keboola Demo" client.close() def test_verify_token_401_error(self, httpx_mock) -> None: diff --git a/tests/test_manage_client.py b/tests/test_manage_client.py index 0e50350f..b6d11d84 100644 --- a/tests/test_manage_client.py +++ b/tests/test_manage_client.py @@ -296,6 +296,25 @@ def test_expires_in_none_excluded_from_payload(self, httpx_mock) -> None: client.close() +class TestGetOrganization: + """Tests for get_organization().""" + + def test_returns_organization_details(self, httpx_mock) -> None: + """get_organization returns id + name (used by org setup to populate org_name).""" + httpx_mock.add_response( + url=f"{STACK_URL}/manage/organizations/438", + json={"id": 438, "name": "Keboola Demo"}, + status_code=200, + ) + + client = ManageClient(stack_url=STACK_URL, manage_token=MANAGE_TOKEN) + result = client.get_organization(438) + + assert result["id"] == 438 + assert result["name"] == "Keboola Demo" + client.close() + + class TestGetProject: """Tests for get_project().""" diff --git a/tests/test_models.py b/tests/test_models.py index 3a12a00e..aed49795 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -38,6 +38,34 @@ def test_default_values(self) -> None: ) assert config.project_name == "" assert config.project_id is None + assert config.org_id is None + assert config.org_name is None + + def test_org_fields_persisted(self) -> None: + """Organization fields round-trip through JSON serialization.""" + config = ProjectConfig( + stack_url="https://connection.keboola.com", + token="901-token", + org_id=438, + org_name="Keboola Demo", + ) + restored = ProjectConfig.model_validate_json(config.model_dump_json()) + assert restored.org_id == 438 + assert restored.org_name == "Keboola Demo" + + def test_legacy_config_without_org_fields(self) -> None: + """Configs persisted before org fields existed still load cleanly.""" + legacy = json.dumps( + { + "stack_url": "https://connection.keboola.com", + "token": "901-token", + "project_name": "Legacy", + "project_id": 100, + } + ) + restored = ProjectConfig.model_validate_json(legacy) + assert restored.org_id is None + assert restored.org_name is None def test_json_round_trip(self) -> None: """ProjectConfig can be serialized to JSON and deserialized back.""" diff --git a/tests/test_org_service.py b/tests/test_org_service.py index caeeaede..8c1f2caf 100644 --- a/tests/test_org_service.py +++ b/tests/test_org_service.py @@ -172,6 +172,70 @@ def storage_factory(url, token): assert "alpha" in config.projects assert "beta" in config.projects + def test_populates_org_name_from_per_project_payload(self, tmp_path: Path) -> None: + """When list_organization_projects payload carries organization.name, use it.""" + config_dir = tmp_path / "config" + config_dir.mkdir() + store = ConfigStore(config_dir=config_dir) + + projects = [ + { + "id": 100, + "name": "Alpha", + "organization": {"id": 42, "name": "Acme Corp"}, + }, + ] + + service = OrgService( + config_store=store, + manage_client_factory=_make_manage_client(projects), + storage_client_factory=_make_storage_client(project_name="Alpha", project_id=100), + ) + + service.setup_organization( + stack_url="https://connection.keboola.com", + manage_token="manage-token-123456789012345678", + org_id=42, + ) + + registered = store.load().projects["alpha"] + assert registered.org_id == 42 + assert registered.org_name == "Acme Corp" + + def test_falls_back_to_get_organization_for_name(self, tmp_path: Path) -> None: + """When per-project payload omits org name, fetch it via get_organization.""" + config_dir = tmp_path / "config" + config_dir.mkdir() + store = ConfigStore(config_dir=config_dir) + + projects = [{"id": 100, "name": "Alpha"}] # no organization sub-object + + manage_mock = MagicMock() + manage_mock.list_organization_projects.return_value = projects + manage_mock.create_project_token.return_value = { + "id": "tok-1", + "token": "901-99999-generatedToken1234567890ab", + "description": "kbagent-cli", + } + manage_mock.get_organization.return_value = {"id": 42, "name": "Resolved Org"} + + service = OrgService( + config_store=store, + manage_client_factory=lambda url, token: manage_mock, + storage_client_factory=_make_storage_client(project_name="Alpha", project_id=100), + ) + + service.setup_organization( + stack_url="https://connection.keboola.com", + manage_token="manage-token-123456789012345678", + org_id=42, + ) + + manage_mock.get_organization.assert_called_once_with(42) + registered = store.load().projects["alpha"] + assert registered.org_id == 42 + assert registered.org_name == "Resolved Org" + def test_skip_existing_projects(self, tmp_path: Path) -> None: """Already-registered projects are skipped.""" config_dir = tmp_path / "config" diff --git a/tests/test_services.py b/tests/test_services.py index 1a802c2e..9c939ea4 100644 --- a/tests/test_services.py +++ b/tests/test_services.py @@ -293,6 +293,44 @@ def test_list_multiple_projects(self, tmp_config_dir: Path) -> None: prod = next(p for p in result if p["alias"] == "prod") assert prod["is_default"] is True + def test_list_projects_includes_org_info(self, tmp_config_dir: Path) -> None: + """When verify_token returns organization info, it's propagated to list_projects.""" + store = ConfigStore(config_dir=tmp_config_dir) + mock_client = make_mock_client(org_id=438, org_name="Keboola Demo") + service = ProjectService( + config_store=store, + client_factory=lambda url, token: mock_client, + ) + + service.add_project( + alias="demo", + stack_url="https://connection.keboola.com", + token="901-10493007-VDtlEDWDF6Tx5V8jjE8FshFlqM0Hl0c08KHqpt0k", + ) + + result = service.list_projects() + assert result[0]["org_id"] == 438 + assert result[0]["org_name"] == "Keboola Demo" + + def test_list_projects_org_missing_yields_none(self, tmp_config_dir: Path) -> None: + """Projects added without org info expose org_id/org_name as None (UI fallback to '—').""" + store = ConfigStore(config_dir=tmp_config_dir) + mock_client = make_mock_client() # no org info + service = ProjectService( + config_store=store, + client_factory=lambda url, token: mock_client, + ) + + service.add_project( + alias="prod", + stack_url="https://connection.keboola.com", + token="901-10493007-VDtlEDWDF6Tx5V8jjE8FshFlqM0Hl0c08KHqpt0k", + ) + + result = service.list_projects() + assert result[0]["org_id"] is None + assert result[0]["org_name"] is None + def test_list_projects_token_never_fully_shown(self, tmp_config_dir: Path) -> None: """list_projects never returns the full token.""" store = ConfigStore(config_dir=tmp_config_dir) diff --git a/web/frontend/src/pages/Projects.tsx b/web/frontend/src/pages/Projects.tsx index 096d431d..1c1f4135 100644 --- a/web/frontend/src/pages/Projects.tsx +++ b/web/frontend/src/pages/Projects.tsx @@ -89,6 +89,17 @@ export function ProjectsPage() { ), }, + { + header: "Org", + cell: (p) => + p.org_name ? ( + {p.org_name} + ) : ( + + — + + ), + }, { header: "Project", cell: (p) => {p.project_name} }, { header: "ID", diff --git a/web/frontend/src/types.ts b/web/frontend/src/types.ts index 564156ff..9410e32d 100644 --- a/web/frontend/src/types.ts +++ b/web/frontend/src/types.ts @@ -12,6 +12,8 @@ export interface Project { token: string; // already masked is_default: boolean; active_branch_id: number | null; + org_id: number | null; + org_name: string | null; } export interface ProjectStatus { From 3d55de256c6a04d6ffbba42d654e46c5ed0c9e66 Mon Sep 17 00:00:00 2001 From: Petr Date: Fri, 15 May 2026 09:15:25 +0200 Subject: [PATCH 03/12] feat(ui): quick search + wider project picker dropdown The top-bar project switcher was unusable once a user had ~6+ projects registered (Vojta hit this with the org-wide setup): no way to filter, and the 288px column truncated longer project names. Changes: - Widen dropdown from w-72 (288px) to w-96 (384px) and bump the max height to 28rem so more items are visible at once. - Show a search input at the top when there are more than 5 projects; filters case-insensitively across alias, project_name, and org_name. - Auto-focus the search field when the menu opens; Escape clears the query (or closes the menu when already empty); Enter picks the only remaining match when the filter narrowed the list to one. - Display org_name (added in #290) as a subtle suffix on each row so users with multi-org setups can disambiguate at a glance. --- web/frontend/src/layout/TopBar.tsx | 118 ++++++++++++++++++++++------- 1 file changed, 91 insertions(+), 27 deletions(-) diff --git a/web/frontend/src/layout/TopBar.tsx b/web/frontend/src/layout/TopBar.tsx index a54704b4..b24c90e2 100644 --- a/web/frontend/src/layout/TopBar.tsx +++ b/web/frontend/src/layout/TopBar.tsx @@ -1,6 +1,6 @@ import { useQuery } from "@tanstack/react-query"; -import { ChevronDown, GitBranch, Layers, Moon, Server, Sun } from "lucide-react"; -import { useEffect, useRef, useState } from "react"; +import { ChevronDown, GitBranch, Layers, Moon, Search, Server, Sun } from "lucide-react"; +import { useEffect, useMemo, useRef, useState } from "react"; import { api } from "../api/client"; import { useUIState } from "../state"; import { useTheme } from "../theme"; @@ -96,7 +96,10 @@ function ProjectPicker({ onChange: (p: string) => void; }) { const [open, setOpen] = useState(false); + const [query, setQuery] = useState(""); const ref = useRef(null); + const searchRef = useRef(null); + useEffect(() => { if (!open) return; const onClick = (e: MouseEvent) => { @@ -105,6 +108,27 @@ function ProjectPicker({ window.addEventListener("mousedown", onClick); return () => window.removeEventListener("mousedown", onClick); }, [open]); + + // Reset query when closing and autofocus the search field when opening. + useEffect(() => { + if (open) { + // Defer one tick: the input is mounted in the same render and + // focusing it before paint occasionally drops on Safari/Firefox. + const t = setTimeout(() => searchRef.current?.focus(), 0); + return () => clearTimeout(t); + } + setQuery(""); + }, [open]); + + const filtered = useMemo(() => { + const q = query.trim().toLowerCase(); + if (!q) return projects; + return projects.filter((p) => { + const haystack = `${p.alias} ${p.project_name} ${p.org_name ?? ""}`.toLowerCase(); + return haystack.includes(q); + }); + }, [projects, query]); + return (
{open ? ( -
- {projects.length === 0 ? ( -
- No projects. Add one via the Projects page. -
- ) : ( - projects.map((p) => ( - - )) - )} + placeholder={`filter ${projects.length} projects (alias, name, org)`} + className="flex-1 bg-transparent text-sm focus:outline-none placeholder-zinc-400 dark:placeholder-zinc-600" + /> + {query ? ( + + {filtered.length}/{projects.length} + + ) : null} +
+ ) : null} +
+ {projects.length === 0 ? ( +
+ No projects. Add one via the Projects page. +
+ ) : filtered.length === 0 ? ( +
+ No project matches “{query}”. +
+ ) : ( + filtered.map((p) => ( + + )) + )} +
) : null}
From 70af73146bd511b8dabedf4451a5674f6051dd97 Mon Sep 17 00:00:00 2001 From: Petr Date: Fri, 15 May 2026 09:19:33 +0200 Subject: [PATCH 04/12] fix(lineage): friendly banner when diagram exceeds Mermaid size limit (#289) Real lineage graphs (multi-org Keboola setups) routinely trip Mermaid's internal source-text guard, which surfaced as a raw 'Maximum text size in diagram exceeded' error message. The diagram silently never rendered and there was no hint about what to do. Replace the raw error with a styled amber banner that: - States plainly that the diagram is too large for the embedded preview - Points to `kbagent lineage server --load ` as the dedicated tool with zoom, pan, search, and column-level drill-down - Reminds the user that the sidebar's upstream/downstream/columns toggles can narrow the edge set back under the limit The banner only fires for the size-limit case; other Mermaid render errors (parse failures, malformed edge data, ...) still show the raw message for debugging. --- web/frontend/src/pages/Lineage.tsx | 37 ++++++++++++++++++++++++++++-- 1 file changed, 35 insertions(+), 2 deletions(-) diff --git a/web/frontend/src/pages/Lineage.tsx b/web/frontend/src/pages/Lineage.tsx index 3f6c5e47..47fef73e 100644 --- a/web/frontend/src/pages/Lineage.tsx +++ b/web/frontend/src/pages/Lineage.tsx @@ -528,13 +528,46 @@ function MermaidGraph({ edges }: { edges: LineageEdge[] }) { ); } + // Mermaid trips its hardcoded text-size guard around ~50 KB of source. + // The error message is "Maximum text size in diagram exceeded"; we also + // accept a generic substring match in case the message wording shifts. + const isOversize = !!error && /maximum text size/i.test(error); + return (

Diagram

- {error ? ( + {isOversize ? ( + + ) : error ? (
{error}
) : null} -
+ {!isOversize ? ( +
+ ) : null} +
+ ); +} + +function OversizeBanner({ edgeCount }: { edgeCount: number }) { + return ( +
+
+ Diagram too large to render here ({edgeCount} edges) +
+

+ Mermaid's embedded renderer caps source-text size and rejected this graph. + The embedded preview is intentionally lightweight — for real exploration + of a large lineage, open the dedicated lineage server which supports + zoom, pan, search, and column-level drill-down: +

+ 
+        kbagent lineage server --load <path-to-lineage.json>
+
+

+ Tip: the upstream / downstream / column toggles in the sidebar narrow + the edge set first — that often brings the diagram back under the + embedded renderer's limit. +

); } From 50a44a152ec3af50b9071d0b028e7b9454233232 Mon Sep 17 00:00:00 2001 From: Petr Date: Fri, 15 May 2026 10:44:28 +0200 Subject: [PATCH 05/12] feat(workspaces): AI-assisted SQL writer in the workspace editor (#287) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds a 'Help me write this SQL' button next to the workspace SQL editor that spawns a local AI CLI (claude / codex / gemini), feeds it a meta-prompt grounded in the user's current workspace, and streams the generated SQL back into the editor. Same mechanism as the AI prompt helper for scheduled agents (#283), so users get a consistent 'local AI does the work' workflow across the UI. The meta-prompt bakes in: - Project alias, backend (snowflake / bigquery / other), default schema - Backend-specific INFORMATION_SCHEMA recipes — claude routinely invents Snowflake-style queries when the workspace is BigQuery unless told the exact dataset-path syntax - The bucket list already loaded in the editor's Storage Explorer (truncated to 50 to keep the prompt tractable) - Real kbagent CLI commands for further discovery (kbagent workspace query, kbagent storage table-detail) so the AI doesn't fabricate flag-less variants - Output contract that forbids ```sql fences and 'Here's the SQL:' preambles, with post-hoc cleanup as a belt-and-suspenders Backend: - build_sql_helper_meta_prompt / clean_sql_helper_response in agent_runner.py — pure-text helpers, fully tested - POST /workspaces/sql/improve/stream — SSE endpoint mirroring /agents/prompt/improve/stream protocol (init/stdout/stderr/done) with a cleaned 'sql' field on the done event Frontend: - SqlHelperPanel component in Workspaces.tsx, mirroring PromptHelperPanel from Agents.tsx (CLI selector, goal textarea, live streaming preview, Apply/Discard) - Wired into SqlEditorDrawer; the AI CLI choice is held by the drawer so it survives toggling the helper open/closed - The editor's draft SQL is passed to the AI as a starting point; Apply replaces the editor body with the generated query Tests: 23 new (12 build/clean unit, 3 SSE endpoint integration). --- src/keboola_agent_cli/server/agent_runner.py | 172 ++++++++ .../server/routers/workspaces.py | 109 ++++- tests/test_workspace_sql_helper.py | 371 ++++++++++++++++++ web/frontend/src/pages/Workspaces.tsx | 353 ++++++++++++++++- 4 files changed, 1002 insertions(+), 3 deletions(-) create mode 100644 tests/test_workspace_sql_helper.py diff --git a/src/keboola_agent_cli/server/agent_runner.py b/src/keboola_agent_cli/server/agent_runner.py index 75fe648f..2a2bbe02 100644 --- a/src/keboola_agent_cli/server/agent_runner.py +++ b/src/keboola_agent_cli/server/agent_runner.py @@ -204,6 +204,178 @@ def build_prompt_helper_meta_prompt( "prompt:", ) +# Same idea for the SQL helper; the AI is told to emit only SQL but routinely +# starts with "Here's the SQL:" or wraps the body in ```sql fences. +_SQL_RESPONSE_PREAMBLES = ( + "here is the sql:", + "here's the sql:", + "here is a sql:", + "here's a sql:", + "here is the query:", + "here's the query:", + "sql:", + "query:", +) + + +def build_sql_helper_meta_prompt( + *, + goal: str, + project: str, + backend: str, + schema: str, + draft_sql: str = "", + bucket_ids: list[str] | None = None, + serve_url: str | None = None, +) -> str: + """Compose the meta-prompt sent to the AI CLI by the workspace SQL helper. + + The AI is asked to produce a single polished SQL statement (or a small + statement batch) that runs against the user's Keboola workspace. It is + explicitly instructed to discover table / column shape via + INFORMATION_SCHEMA using the kbagent CLI before guessing column names. + + Backend-specific hints are folded in so claude doesn't have to "know" the + quirks: BigQuery's backticked dataset paths and per-dataset + INFORMATION_SCHEMA, Snowflake's CURRENT_SCHEMA() default, etc. The + bucket list (when supplied) gives the AI a starting catalog without + burning a tool call. + """ + goal_clean = goal.strip() + draft_block = ( + f"USER'S CURRENT DRAFT (refine this, don't throw it away):\n{draft_sql.strip()}" + if draft_sql.strip() + else "USER'S CURRENT DRAFT: (empty -- write the query from scratch.)" + ) + bucket_block = ( + "VISIBLE BUCKETS (already loaded in the editor sidebar):\n" + + "\n".join(f" - {b}" for b in bucket_ids[:50]) + if bucket_ids + else "VISIBLE BUCKETS: (none preloaded -- discover via INFORMATION_SCHEMA.)" + ) + if len(bucket_ids or []) > 50: + bucket_block += f"\n ... and {len(bucket_ids or []) - 50} more (truncated)" + + backend_hint = _sql_helper_backend_hint(backend, schema) + serve_hint = ( + f"SERVE CONTEXT: kbagent serve is reachable at {serve_url}; the AI agent\n" + "shell has KBAGENT_SERVE_URL + KBAGENT_SERVE_TOKEN env vars pre-set, so\n" + "`kbagent http get /...` is the fastest discovery path." + if serve_url + else "SERVE CONTEXT: assume `kbagent` CLI is available on PATH." + ) + + return f"""\ +You are a senior data engineer writing SQL for a Keboola workspace. Your +output will be pasted into the workspace SQL editor verbatim and executed +through the Keboola Query Service against project '{project}'. The Query +Service runs SELECT only -- it rejects SHOW / DESCRIBE / DDL / DML. + +WORKSPACE CONTEXT: +- Project alias: {project} +- Backend: {backend} +- Default schema: {schema} + +USER'S GOAL (plain English): +{goal_clean} + +{draft_block} + +{bucket_block} + +{backend_hint} + +DISCOVERY (do this BEFORE guessing column names): +- Use `kbagent workspace query --project {project} --workspace-id --sql '...'` + with an INFORMATION_SCHEMA query to confirm table + column names exist. +- Alternative: `kbagent storage table-detail --project {project} --table-id ` + returns the full column list for a Storage table without spinning up a query. +{serve_hint} + +REQUIREMENTS for the returned SQL: +- Match the user's goal precisely; do not invent columns. +- Be a single SELECT statement (or a tiny CTE batch) -- nothing destructive. +- Qualify tables explicitly when joining across buckets so the result is + unambiguous after the workspace is reused. +- Add a brief 1-line `-- comment` at the top describing what the query + returns (purpose + key filters), but no other prose. + +OUTPUT CONTRACT (critical): +- Output ONLY the SQL. Plain text. +- Do NOT wrap the SQL in ```sql fences. +- Do NOT prefix with "Here's the SQL:" / "Rewritten query:" / similar. +- Do NOT append commentary after the SQL. +""" + + +def _sql_helper_backend_hint(backend: str, schema: str) -> str: + """Emit backend-specific INFORMATION_SCHEMA recipes for the meta-prompt. + + Keboola Workspaces run on three backends; each has different table-catalog + surface area, so the meta-prompt embeds the exact INFORMATION_SCHEMA query + the AI should run for discovery. Without this hint claude routinely + invents Snowflake-style queries when the workspace is BigQuery. + """ + backend_lc = (backend or "").lower() + if backend_lc == "bigquery": + return ( + "BACKEND HINT (BigQuery):\n" + f"- Workspace schema is the dataset `{schema}`.\n" + f"- Backtick-quote dataset + table names: `\\`{schema}\\`.\\`\\``.\n" + f"- Discovery: SELECT table_name FROM `{schema}.INFORMATION_SCHEMA.TABLES`;\n" + f"- Columns: SELECT column_name, data_type FROM " + f"`{schema}.INFORMATION_SCHEMA.COLUMNS` WHERE table_name='
';" + ) + if backend_lc == "snowflake": + return ( + "BACKEND HINT (Snowflake):\n" + f"- Workspace default schema is `{schema}`. Identifiers are case-sensitive\n" + f' when quoted; Keboola Storage tables are quoted ("my-table").\n' + f"- Discovery: SELECT TABLE_NAME, ROW_COUNT FROM INFORMATION_SCHEMA.TABLES\n" + f" WHERE TABLE_SCHEMA = CURRENT_SCHEMA();\n" + f"- Columns: SELECT COLUMN_NAME, DATA_TYPE FROM INFORMATION_SCHEMA.COLUMNS\n" + f" WHERE TABLE_SCHEMA = CURRENT_SCHEMA() AND TABLE_NAME = '
';" + ) + # Unknown / future backend: stay generic so the AI still has a starting point. + return ( + f"BACKEND HINT ({backend or 'unknown'}):\n" + f"- Workspace default schema is `{schema}`.\n" + "- Discovery: query INFORMATION_SCHEMA.TABLES / COLUMNS following the\n" + " backend's conventions (Snowflake = CURRENT_SCHEMA(), BigQuery = dataset\n" + " path, Postgres = current_schema())." + ) + + +def clean_sql_helper_response(text: str) -> str: + """Strip code fences, preambles, and claude jsonl duplication from SQL output. + + Mirrors :func:`clean_prompt_helper_response` step-for-step but uses the + SQL-specific preamble list. Two distinct cleaners (instead of a unified + one with a knob) makes the call sites self-documenting and lets future + SQL/prompt divergence land without entangling. + """ + text = text.strip() + # Step 1: collapse "AB" where A == B (claude jsonl duplication). + if text and len(text) % 2 == 0: + half = len(text) // 2 + if text[:half] == text[half:]: + text = text[:half].rstrip() + # Step 2: strip a single set of leading/trailing code fences. Accept + # ```sql or ``` -- the AI uses both interchangeably. + if text.startswith("```"): + nl = text.find("\n") + if nl != -1: + text = text[nl + 1 :] + if text.endswith("```"): + text = text[:-3] + text = text.strip() + # Step 3: strip a preamble like "Here's the SQL:\n\n..." on the first line. + lines = text.split("\n", 1) + first = lines[0].strip().lower() + if any(first == p or first.startswith(p) for p in _SQL_RESPONSE_PREAMBLES): + text = lines[1].strip() if len(lines) > 1 else "" + return text.strip() + def clean_prompt_helper_response(text: str) -> str: """Trim surrounding code fences, preambles, and dedup the response. diff --git a/src/keboola_agent_cli/server/routers/workspaces.py b/src/keboola_agent_cli/server/routers/workspaces.py index 91c4fec6..003839ad 100644 --- a/src/keboola_agent_cli/server/routers/workspaces.py +++ b/src/keboola_agent_cli/server/routers/workspaces.py @@ -2,9 +2,12 @@ from __future__ import annotations +import json +from collections.abc import AsyncIterator from typing import Any -from fastapi import APIRouter, Depends, Query +from fastapi import APIRouter, Depends, HTTPException, Query +from fastapi.responses import StreamingResponse from pydantic import BaseModel from ..dependencies import ServiceRegistry, get_registry @@ -35,6 +38,26 @@ class FromTransformation(BaseModel): row_id: str | None = None +class SqlHelperRequest(BaseModel): + """Input for the /workspaces/sql/improve/stream endpoint. + + Mirrors :class:`PromptHelperRequest` from the agents router but carries + workspace-specific context (project, backend, schema, visible buckets) + so the meta-prompt the AI receives is grounded in the user's current + workspace -- no generic 'write SQL' guesswork. + """ + + cli: str # claude | codex | gemini -- same recipe as ai_agent runs + goal: str + project: str + backend: str + schema_name: str + workspace_id: int | None = None + draft_sql: str = "" + bucket_ids: list[str] = [] + extra_args: list[str] = [] + + @router.get("") def list_workspaces( project: list[str] | None = Query(None), @@ -108,6 +131,90 @@ def query( ) +def _sse(event: str, data: dict[str, Any]) -> bytes: + """Encode a single SSE frame (event + data).""" + return f"event: {event}\ndata: {json.dumps(data, default=str)}\n\n".encode() + + +@router.post("/sql/improve/stream") +async def improve_sql_stream( + body: SqlHelperRequest, + registry: ServiceRegistry = Depends(get_registry), +) -> StreamingResponse: + """Stream an AI-generated SQL query back to the workspace SQL editor. + + Mirrors /agents/prompt/improve/stream but with a SQL-specific meta-prompt + that grounds the AI in the workspace's backend (snowflake/bigquery), + default schema, and the visible bucket catalog the editor sidebar has + already loaded. The AI is also told how to use INFORMATION_SCHEMA via + `kbagent workspace query` for any discovery the bucket hint doesn't cover. + + Same SSE event protocol as the agent prompt helper (init/stdout/stderr/ + done) so the UI can reuse the streaming progress renderer. + """ + from ..agent_runner import ( + build_sql_helper_meta_prompt, + clean_sql_helper_response, + stream_ai_agent_events, + ) + + goal = body.goal.strip() + if not goal: + raise HTTPException(status_code=400, detail="goal must not be empty") + + meta_prompt = build_sql_helper_meta_prompt( + goal=goal, + project=body.project, + backend=body.backend, + schema=body.schema_name, + draft_sql=body.draft_sql, + bucket_ids=body.bucket_ids or None, + ) + params: dict[str, Any] = { + "cli": body.cli, + "prompt": meta_prompt, + "extra_args": body.extra_args, + # SQL helper prompts target ~10-30s for a simple SELECT, up to a minute + # when the AI has to round-trip INFORMATION_SCHEMA. 180s cap matches + # the prompt helper so a stuck CLI doesn't camp on the connection. + "timeout": 180.0, + } + + async def gen() -> AsyncIterator[bytes]: + yield _sse( + "init", + { + "kind": "sql_helper", + "cli": body.cli, + "project": body.project, + "backend": body.backend, + "goal_preview": goal[:200], + }, + ) + try: + async for evt in stream_ai_agent_events(registry, params): + if evt["event"] == "done": + raw = str(evt["data"].get("response") or "") + cleaned = clean_sql_helper_response(raw) + enriched = {**evt["data"], "sql": cleaned, "raw_response": raw} + yield _sse("done", enriched) + else: + yield _sse(evt["event"], evt["data"]) + except ValueError as exc: + yield _sse("done", {"status": "error", "error": str(exc)}) + except Exception as exc: + yield _sse("done", {"status": "error", "error": str(exc)}) + + return StreamingResponse( + gen(), + media_type="text/event-stream", + headers={ + "Cache-Control": "no-cache, no-transform", + "X-Accel-Buffering": "no", + }, + ) + + @router.post("/{project}/from-transformation") def from_transformation( project: str, diff --git a/tests/test_workspace_sql_helper.py b/tests/test_workspace_sql_helper.py new file mode 100644 index 00000000..925b58a6 --- /dev/null +++ b/tests/test_workspace_sql_helper.py @@ -0,0 +1,371 @@ +"""Tests for the AI-driven SQL helper used by the workspace SQL editor. + +Mirrors :mod:`tests.test_agent_prompt_helper` but for the workspace-side +helpers: + +- ``build_sql_helper_meta_prompt`` / ``clean_sql_helper_response`` -- pure + text helpers in :mod:`keboola_agent_cli.server.agent_runner`. +- ``POST /workspaces/sql/improve/stream`` -- the SSE endpoint that wires + the helpers into the chosen AI CLI. We mock :func:`stream_ai_agent_events` + so the test does not spawn a real ``claude`` / ``codex`` / ``gemini`` + subprocess and verify the endpoint forwards events and enriches the + ``done`` payload with a cleaned ``sql`` field. +""" + +from __future__ import annotations + +import importlib.util +from pathlib import Path + +import pytest + +if importlib.util.find_spec("fastapi") is None: # pragma: no cover + pytest.skip( + "FastAPI not installed; run `uv pip install -e '.[server]'`", + allow_module_level=True, + ) + +from fastapi.testclient import TestClient + +from keboola_agent_cli.server import create_app +from keboola_agent_cli.server.agent_runner import ( + build_sql_helper_meta_prompt, + clean_sql_helper_response, +) + +# --------------------------------------------------------------------- +# build_sql_helper_meta_prompt +# --------------------------------------------------------------------- + + +class TestBuildSqlHelperMetaPrompt: + def test_includes_goal_verbatim(self) -> None: + prompt = build_sql_helper_meta_prompt( + goal="top 10 orders by revenue last 30 days", + project="demo", + backend="snowflake", + schema="WORKSPACE_123", + ) + assert "top 10 orders by revenue last 30 days" in prompt + + def test_includes_project_and_backend(self) -> None: + prompt = build_sql_helper_meta_prompt( + goal="g", + project="padak", + backend="bigquery", + schema="my_dataset", + ) + assert "padak" in prompt + assert "bigquery" in prompt + assert "my_dataset" in prompt + + def test_draft_block_present_when_draft_given(self) -> None: + prompt = build_sql_helper_meta_prompt( + goal="g", + project="demo", + backend="snowflake", + schema="s", + draft_sql="SELECT * FROM orders LIMIT 5", + ) + assert "SELECT * FROM orders LIMIT 5" in prompt + assert "refine this" in prompt + + def test_empty_draft_marked_explicitly(self) -> None: + prompt = build_sql_helper_meta_prompt( + goal="g", + project="demo", + backend="snowflake", + schema="s", + draft_sql=" ", + ) + assert "empty -- write the query from scratch" in prompt + + def test_bucket_hint_lists_visible_buckets(self) -> None: + prompt = build_sql_helper_meta_prompt( + goal="g", + project="demo", + backend="snowflake", + schema="s", + bucket_ids=["in.c-orders", "in.c-customers"], + ) + assert "in.c-orders" in prompt + assert "in.c-customers" in prompt + assert "VISIBLE BUCKETS" in prompt + + def test_bucket_hint_truncates_at_50(self) -> None: + """Long bucket lists are truncated so the prompt stays tractable.""" + prompt = build_sql_helper_meta_prompt( + goal="g", + project="demo", + backend="snowflake", + schema="s", + bucket_ids=[f"in.c-b{i}" for i in range(75)], + ) + assert "and 25 more (truncated)" in prompt + + def test_bigquery_backend_hint_specifies_dataset_path(self) -> None: + prompt = build_sql_helper_meta_prompt( + goal="g", + project="demo", + backend="BigQuery", # mixed case to test normalization + schema="my_dataset", + ) + assert "BACKEND HINT (BigQuery)" in prompt + assert "`my_dataset.INFORMATION_SCHEMA.TABLES`" in prompt + + def test_snowflake_backend_hint_uses_current_schema(self) -> None: + prompt = build_sql_helper_meta_prompt( + goal="g", + project="demo", + backend="snowflake", + schema="WORKSPACE_X", + ) + assert "BACKEND HINT (Snowflake)" in prompt + assert "CURRENT_SCHEMA()" in prompt + + def test_unknown_backend_gets_generic_hint(self) -> None: + prompt = build_sql_helper_meta_prompt( + goal="g", + project="demo", + backend="postgres", + schema="s", + ) + assert "BACKEND HINT (postgres)" in prompt + # Generic hint mentions both Snowflake and BigQuery conventions. + assert "current_schema()" in prompt + + def test_output_contract_present(self) -> None: + """The OUTPUT CONTRACT block is the whole reason the helper works. + + Without it, claude wraps SQL in ```sql fences or starts with + "Here's the SQL:" and the user's editor receives garbage. Pin the + contract so accidental trimming surfaces as a test failure. + """ + prompt = build_sql_helper_meta_prompt( + goal="g", + project="demo", + backend="snowflake", + schema="s", + ) + assert "OUTPUT CONTRACT" in prompt + assert "Output ONLY the SQL" in prompt + assert "```sql" in prompt # the contract names the forbidden fence form + + def test_discovery_block_references_kbagent_commands(self) -> None: + """The meta-prompt anchors the AI in REAL kbagent commands for + catalog discovery so it doesn't invent flag-less variants. + """ + prompt = build_sql_helper_meta_prompt( + goal="g", + project="demo", + backend="snowflake", + schema="s", + ) + assert "kbagent workspace query" in prompt + assert "kbagent storage table-detail" in prompt + + +# --------------------------------------------------------------------- +# clean_sql_helper_response +# --------------------------------------------------------------------- + + +class TestCleanSqlHelperResponse: + def test_passthrough_clean_sql(self) -> None: + body = ( + "-- top orders by revenue\nSELECT id, total FROM orders ORDER BY total DESC LIMIT 10;" + ) + assert clean_sql_helper_response(body) == body + + def test_strips_sql_code_fence(self) -> None: + raw = "```sql\nSELECT 1;\n```" + assert clean_sql_helper_response(raw) == "SELECT 1;" + + def test_strips_plain_code_fence(self) -> None: + raw = "```\nSELECT 1;\n```" + assert clean_sql_helper_response(raw) == "SELECT 1;" + + def test_strips_here_is_the_sql_preamble(self) -> None: + raw = "Here's the SQL:\n\nSELECT 1;" + assert clean_sql_helper_response(raw) == "SELECT 1;" + + def test_strips_query_preamble(self) -> None: + raw = "Query:\nSELECT 1;" + assert clean_sql_helper_response(raw) == "SELECT 1;" + + def test_strips_surrounding_whitespace(self) -> None: + assert clean_sql_helper_response("\n\n SELECT 1; \n\n") == "SELECT 1;" + + def test_empty_input_yields_empty(self) -> None: + assert clean_sql_helper_response("") == "" + + def test_dedups_exact_double_body(self) -> None: + """Same claude jsonl duplication quirk as the prompt helper.""" + body = "SELECT id FROM orders ORDER BY total DESC LIMIT 10;" + raw = body + body + assert clean_sql_helper_response(raw) == body + + def test_does_not_dedup_unrelated_halves(self) -> None: + raw = "SELECT a FROM t1;\nSELECT b FROM t2;" + # The two halves aren't equal -- the function must leave it alone. + assert clean_sql_helper_response(raw) == raw + + +# --------------------------------------------------------------------- +# POST /workspaces/sql/improve/stream +# --------------------------------------------------------------------- + + +@pytest.fixture +def client(tmp_path: Path) -> TestClient: + app = create_app(config_dir=str(tmp_path), auth_token="test-token") + return TestClient(app) + + +def _parse_sse_events(text: str) -> list[tuple[str, str]]: + """Parse the SSE wire-format body into [(event, data_json_string)].""" + out: list[tuple[str, str]] = [] + event = "message" + data = "" + for line in text.splitlines(): + if line == "": + if data: + out.append((event, data)) + event = "message" + data = "" + continue + if line.startswith(":"): + continue + if line.startswith("event:"): + event = line[6:].strip() + elif line.startswith("data:"): + data = line[5:].lstrip() + if data: + out.append((event, data)) + return out + + +class TestImproveSqlStreamEndpoint: + def test_empty_goal_rejected_400(self, client: TestClient) -> None: + res = client.post( + "/workspaces/sql/improve/stream", + json={ + "cli": "claude", + "goal": " ", + "project": "demo", + "backend": "snowflake", + "schema_name": "WORKSPACE_X", + }, + headers={"Authorization": "Bearer test-token"}, + ) + assert res.status_code == 400, res.text + + def test_streams_events_and_enriches_done_with_sql( + self, + client: TestClient, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """Happy path: mocked AI stream events get forwarded, and the final + ``done`` payload carries a cleaned ``sql`` field plus the raw + response for debugging. + """ + + async def fake_stream( + registry: object, + params: dict[str, object], + ): + yield {"event": "stdout", "data": {"raw": "thinking..."}} + yield { + "event": "done", + "data": { + "cli": "claude", + "status": "ok", + "exit_code": 0, + "elapsed_seconds": 1.5, + "response": "```sql\nSELECT id FROM orders LIMIT 10;\n```", + "stderr": "", + }, + } + + # The endpoint imports stream_ai_agent_events locally inside the + # handler; patching the module-level name in agent_runner works + # because Python re-resolves it on each call. + monkeypatch.setattr( + "keboola_agent_cli.server.agent_runner.stream_ai_agent_events", + fake_stream, + ) + + with client.stream( + "POST", + "/workspaces/sql/improve/stream", + json={ + "cli": "claude", + "goal": "top 10 orders by id", + "project": "demo", + "backend": "snowflake", + "schema_name": "WORKSPACE_X", + "bucket_ids": ["in.c-orders"], + }, + headers={"Authorization": "Bearer test-token"}, + ) as res: + assert res.status_code == 200 + body = res.read().decode("utf-8") + + events = _parse_sse_events(body) + names = [e for e, _ in events] + assert "init" in names + assert "stdout" in names + assert "done" in names + + import json as _json + + done_data = next(_json.loads(payload) for evt, payload in events if evt == "done") + # ```sql fence must be stripped before the SQL reaches the UI. + assert done_data["sql"] == "SELECT id FROM orders LIMIT 10;" + # raw_response preserved for debugging. + assert "```sql" in done_data["raw_response"] + assert done_data["status"] == "ok" + + def test_stream_error_surfaces_as_done_error( + self, + client: TestClient, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """A ValueError from ``stream_ai_agent_events`` must surface as a + final ``done`` event with status ``error`` so the React side can + render it without hanging on an unterminated stream. + """ + + async def fake_stream( + registry: object, + params: dict[str, object], + ): + raise ValueError("ai_agent.cli must be one of ['claude', 'codex', 'gemini']") + yield # pragma: no cover -- async generator marker + + monkeypatch.setattr( + "keboola_agent_cli.server.agent_runner.stream_ai_agent_events", + fake_stream, + ) + + with client.stream( + "POST", + "/workspaces/sql/improve/stream", + json={ + "cli": "bogus", + "goal": "anything", + "project": "demo", + "backend": "snowflake", + "schema_name": "s", + }, + headers={"Authorization": "Bearer test-token"}, + ) as res: + assert res.status_code == 200 + body = res.read().decode("utf-8") + + events = _parse_sse_events(body) + import json as _json + + done_data = next(_json.loads(payload) for evt, payload in events if evt == "done") + assert done_data["status"] == "error" + assert "ai_agent.cli" in done_data["error"] diff --git a/web/frontend/src/pages/Workspaces.tsx b/web/frontend/src/pages/Workspaces.tsx index 1aa365f1..015c4fde 100644 --- a/web/frontend/src/pages/Workspaces.tsx +++ b/web/frontend/src/pages/Workspaces.tsx @@ -12,15 +12,32 @@ import { Sparkles, Trash2, Upload, + X, } from "lucide-react"; -import { useState } from "react"; -import { api } from "../api/client"; +import { useEffect, useRef, useState } from "react"; +import { api, ssePost, type SsePostHandle } from "../api/client"; import { Drawer } from "../components/Drawer"; import { Empty, ErrorBox, Loading, PageTitle, TwoPathEmpty } from "../components/Empty"; import { DataTable } from "../components/Table"; import { useUIState } from "../state"; import type { ProjectError, Workspace } from "../types"; +/** + * AbortError shape detection across browsers (DOMException on standards, + * named Error on some shims). Kept inline so a future move of this util + * to a shared module doesn't fan out to every page that needs it. + */ +function isAbortError(err: unknown): boolean { + if (err instanceof DOMException && err.name === "AbortError") return true; + if (err instanceof Error && err.name === "AbortError") return true; + return Boolean( + err && + typeof err === "object" && + "message" in err && + String((err as { message: unknown }).message).toLowerCase().includes("abort"), + ); +} + interface WorkspacesResp { workspaces: Workspace[]; errors: ProjectError[]; @@ -457,6 +474,9 @@ SELECT current_timestamp() AS now;`); const [result, setResult] = useState(null); const [error, setError] = useState(null); const [hint, setHint] = useState(null); + // The AI CLI choice is persisted on the SqlEditorDrawer (not the panel) + // so the user's pick survives toggling the helper open/closed. + const [aiCli, setAiCli] = useState<"claude" | "codex" | "gemini">("claude"); // Fetch buckets + tables from the workspace's project so users can click // them into the editor (Storage Explorer pattern from Keboola UI). @@ -575,6 +595,17 @@ SELECT current_timestamp() AS now;`); {/* Editor + results */}
+ b.id)} + onApply={(generatedSql) => setSql(generatedSql)} + />
); } + +/** + * Inline AI helper for the workspace SQL editor (#287). + * + * Modeled on Agents.tsx > PromptHelperPanel — same SSE protocol (init / + * stdout / stderr / done), same claude / codex / gemini CLI selector, same + * "live preview while streaming → final suggestion → Apply / Discard" + * workflow. The differences: + * + * - Endpoint is /workspaces/sql/improve/stream (workspace-grounded + * meta-prompt: backend, schema, visible buckets are passed in). + * - The `done` payload carries `sql` (not `prompt`); `onApply` replaces + * the editor body with it instead of a prompt textarea. + * - The cancel-on-unmount cleanup is mandatory — without it, the backend + * keeps the claude/codex/gemini subprocess alive while waiting for an + * SSE consumer that will never return. + */ +function SqlHelperPanel({ + cli, + onCliChange, + project, + workspaceId, + backend, + schemaName, + draftSql, + bucketIds, + onApply, +}: { + cli: "claude" | "codex" | "gemini"; + onCliChange: (c: "claude" | "codex" | "gemini") => void; + project: string; + workspaceId: number; + backend: string; + schemaName: string; + draftSql: string; + bucketIds: string[]; + onApply: (sql: string) => void; +}) { + const [open, setOpen] = useState(false); + const [goal, setGoal] = useState(""); + const [running, setRunning] = useState(false); + const [elapsed, setElapsed] = useState(0); + const [livePreview, setLivePreview] = useState(""); + const [finalSql, setFinalSql] = useState(null); + const [error, setError] = useState(null); + const handleRef = useRef(null); + + const reset = () => { + setLivePreview(""); + setFinalSql(null); + setError(null); + setElapsed(0); + }; + + const start = () => { + if (!goal.trim()) { + setError("Describe the query you want first (e.g. 'top 10 customers by revenue last 30 days')."); + return; + } + if (handleRef.current) { + handleRef.current.abort(); + handleRef.current = null; + } + reset(); + setRunning(true); + const startMs = Date.now(); + const tick = setInterval( + () => setElapsed(Math.round((Date.now() - startMs) / 1000)), + 500, + ); + let assistantText = ""; + const handle = ssePost( + "/workspaces/sql/improve/stream", + { + cli, + goal, + project, + backend, + schema_name: schemaName, + workspace_id: workspaceId, + draft_sql: draftSql, + bucket_ids: bucketIds, + }, + { + init: () => { + /* the running indicator already covers this */ + }, + stdout: (d) => { + const data = (d ?? {}) as Record; + // Claude stream-json: assistant turns carry message.content[].text. + if (data.type === "assistant" && typeof data.message === "object") { + const msg = data.message as Record; + const content = msg.content; + if (Array.isArray(content)) { + for (const block of content) { + if ( + block && + typeof block === "object" && + (block as Record).type === "text" && + typeof (block as Record).text === "string" + ) { + assistantText += (block as Record).text as string; + setLivePreview(assistantText); + } + } + } + } else if (typeof data.raw === "string") { + // codex / gemini stream raw text lines (no jsonl). + assistantText += (assistantText ? "\n" : "") + data.raw; + setLivePreview(assistantText); + } + }, + stderr: () => { + /* progress notes — ignored for the preview pane */ + }, + done: (d) => { + const data = (d ?? {}) as Record; + if (data.status === "error") { + setError(String(data.error ?? "AI helper failed")); + return; + } + const cleaned = typeof data.sql === "string" ? data.sql.trim() : ""; + if (!cleaned) { + setError("AI returned an empty query. Refine the goal and regenerate."); + return; + } + setFinalSql(cleaned); + }, + message: () => { + /* unknown event — ignore */ + }, + }, + ); + handleRef.current = handle; + handle.done + .catch((err) => { + if (isAbortError(err)) return; + setError((err as Error).message); + }) + .finally(() => { + clearInterval(tick); + setRunning(false); + handleRef.current = null; + }); + }; + + const cancel = () => { + if (handleRef.current) { + handleRef.current.abort(); + handleRef.current = null; + } + setRunning(false); + }; + + useEffect(() => { + return () => { + if (handleRef.current) { + handleRef.current.abort(); + handleRef.current = null; + } + }; + }, []); + + if (!open) { + return ( +
+ + + uses {cli} with your workspace context (project, backend, visible buckets) baked in + +
+ ); + } + + return ( +
+
+
+ + AI SQL helper · {cli} +
+ +
+ +
+ CLI: + {(["claude", "codex", "gemini"] as const).map((c) => ( + + ))} +
+ +